1 files changed, 41 insertions, 9 deletions
diff --git a/db/migrate/20180812173710_copy_status_stats.rb b/db/migrate/20180812173710_copy_status_stats.rb
index 850aa9c13..ff10c18d9 100644
--- a/db/migrate/20180812173710_copy_status_stats.rb
+++ b/db/migrate/20180812173710_copy_status_stats.rb
@@ -3,15 +3,10 @@ class CopyStatusStats < ActiveRecord::Migration[5.2]
 
   def up
     safety_assured do
-      Status.unscoped.select('id').find_in_batches(batch_size: 5_000) do |statuses|
-        execute <<-SQL.squish
-          INSERT INTO status_stats (status_id, reblogs_count, favourites_count, created_at, updated_at)
-          SELECT id, reblogs_count, favourites_count, created_at, updated_at
-          FROM statuses
-          WHERE id IN (#{statuses.map(&:id).join(', ')})
-          ON CONFLICT (status_id) DO UPDATE
-          SET reblogs_count = EXCLUDED.reblogs_count, favourites_count = EXCLUDED.favourites_count
-        SQL
+      if supports_upsert?
+        up_fast
+      else
+        up_slow
       end
     end
   end
@@ -19,4 +14,41 @@ class CopyStatusStats < ActiveRecord::Migration[5.2]
   def down
     # Nothing
   end
+
+  private
+
+  def supports_upsert?
+    version = select_one("SELECT current_setting('server_version_num') AS v")['v'].to_i
+    version >= 90500
+  end
+
+  def up_fast
+    say 'Upsert is available, importing counters using the fast method'
+
+    Status.unscoped.select('id').find_in_batches(batch_size: 5_000) do |statuses|
+      execute <<-SQL.squish
+        INSERT INTO status_stats (status_id, reblogs_count, favourites_count, created_at, updated_at)
+        SELECT id, reblogs_count, favourites_count, created_at, updated_at
+        FROM statuses
+        WHERE id IN (#{statuses.map(&:id).join(', ')})
+        ON CONFLICT (status_id) DO UPDATE
+        SET reblogs_count = EXCLUDED.reblogs_count, favourites_count = EXCLUDED.favourites_count
+      SQL
+    end
+  end
+
+  def up_slow
+    say 'Upsert is not available in PostgreSQL below 9.5, falling back to slow import of counters'
+
+    # We cannot use bulk INSERT or overarching transactions here because of possible
+    # uniqueness violations that we need to skip over
+    Status.unscoped.select('id, reblogs_count, favourites_count, created_at, updated_at').find_each do |status|
+      begin
+        params = [[nil, status.id], [nil, status.reblogs_count], [nil, status.favourites_count], [nil, status.created_at], [nil, status.updated_at]]
+        exec_insert('INSERT INTO status_stats (status_id, reblogs_count, favourites_count, created_at, updated_at) VALUES ($1, $2, $3, $4, $5)', nil, params)
+      rescue ActiveRecord::RecordNotUnique
+        next
+      end
+    end
+  end
 end