From b2fee68b3f40146e6e482450116fda7236322bd8 Mon Sep 17 00:00:00 2001
From: Gerhard Schlager <gerhard.schlager@discourse.org>
Date: Mon, 7 Aug 2023 12:23:58 +0200
Subject: [PATCH] DEV: Add rake task for generating avatars from SSO

---
 app/models/user_avatar.rb  |  1 +
 app/models/user_profile.rb |  1 +
 lib/tasks/import.rake      | 69 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 71 insertions(+)

diff --git a/app/models/user_avatar.rb b/app/models/user_avatar.rb
index fa830b3fa46..90f396d384d 100644
--- a/app/models/user_avatar.rb
+++ b/app/models/user_avatar.rb
@@ -118,6 +118,7 @@ class UserAvatar < ActiveRecord::Base
         max_file_size: SiteSetting.max_image_size_kb.kilobytes,
         tmp_file_name: "sso-avatar",
         follow_redirect: true,
+        skip_rate_limit: !!options&.fetch(:skip_rate_limit, false),
       )
 
     return unless tempfile
diff --git a/app/models/user_profile.rb b/app/models/user_profile.rb
index dc621885653..91920dd9d70 100644
--- a/app/models/user_profile.rb
+++ b/app/models/user_profile.rb
@@ -117,6 +117,7 @@ class UserProfile < ActiveRecord::Base
         max_file_size: SiteSetting.max_image_size_kb.kilobytes,
         tmp_file_name: "sso-profile-background",
         follow_redirect: true,
+        skip_rate_limit: true,
       )
 
     return unless tempfile
diff --git a/lib/tasks/import.rake b/lib/tasks/import.rake
index 84f4d011529..e8f62e8f4c0 100644
--- a/lib/tasks/import.rake
+++ b/lib/tasks/import.rake
@@ -573,3 +573,72 @@ task "import:update_first_post_created_at" => :environment do
 
   log "Done"
 end
+
+desc "Update avatars from external_avatar_url in SSO records"
+task "import:update_avatars_from_sso" => :environment do
+  log "Updating avatars from SSO records"
+
+  sql = <<~SQL
+    SELECT user_id, external_avatar_url
+    FROM single_sign_on_records s
+    WHERE NOT EXISTS (
+      SELECT 1
+      FROM user_avatars a
+      WHERE a.user_id = s.user_id
+    )
+  SQL
+
+  queue = SizedQueue.new(1000)
+  threads = []
+
+  threads << Thread.new do ||
+    DB.query_each(sql) { |row| queue << { user_id: row.user_id, url: row.external_avatar_url } }
+    queue.close
+  end
+
+  max_count = DB.query_single(<<~SQL).first
+    SELECT COUNT(*)
+    FROM single_sign_on_records s
+    WHERE NOT EXISTS (
+      SELECT 1
+      FROM user_avatars a
+      WHERE a.user_id = s.user_id
+    )
+  SQL
+
+  status_queue = Queue.new
+  status_thread =
+    Thread.new do
+      error_count = 0
+      current_count = 0
+
+      while !(status = status_queue.pop).nil?
+        error_count += 1 if !status
+        current_count += 1
+
+        print "\r%7d / %7d (%d errors)" % [current_count, max_count, error_count]
+      end
+    end
+
+  20.times do
+    threads << Thread.new do
+      while row = queue.pop
+        begin
+          UserAvatar.import_url_for_user(
+            row[:url],
+            User.find(row[:user_id]),
+            override_gravatar: true,
+            skip_rate_limit: true,
+          )
+          status_queue << true
+        rescue StandardError
+          status_queue << false
+        end
+      end
+    end
+  end
+
+  threads.each(&:join)
+  status_queue.close
+  status_thread.join
+end