about summary refs log tree commit diff
path: root/lib/mastodon/search_cli.rb
diff options
context:
space:
mode:
Diffstat (limited to 'lib/mastodon/search_cli.rb')
-rw-r--r--lib/mastodon/search_cli.rb12
1 files changed, 9 insertions, 3 deletions
diff --git a/lib/mastodon/search_cli.rb b/lib/mastodon/search_cli.rb
index 2d1ca1c05..6ad9d7b6a 100644
--- a/lib/mastodon/search_cli.rb
+++ b/lib/mastodon/search_cli.rb
@@ -17,10 +17,11 @@ module Mastodon
     ].freeze
 
     option :concurrency, type: :numeric, default: 2, aliases: [:c], desc: 'Workload will be split between this number of threads'
+    option :batch_size, type: :numeric, default: 1_000, aliases: [:b], desc: 'Number of records in each batch'
     option :only, type: :array, enum: %w(accounts tags statuses), desc: 'Only process these indices'
-    desc 'deploy', 'Create or upgrade ElasticSearch indices and populate them'
+    desc 'deploy', 'Create or upgrade Elasticsearch indices and populate them'
     long_desc <<~LONG_DESC
-      If ElasticSearch is empty, this command will create the necessary indices
+      If Elasticsearch is empty, this command will create the necessary indices
       and then import data from the database into those indices.
 
       This command will also upgrade indices if the underlying schema has been
@@ -35,6 +36,11 @@ module Mastodon
         exit(1)
       end
 
+      if options[:batch_size] < 1
+        say('Cannot run with this batch_size setting, must be at least 1', :red)
+        exit(1)
+      end
+
       indices = begin
         if options[:only]
           options[:only].map { |str| "#{str.camelize}Index".constantize }
@@ -73,7 +79,7 @@ module Mastodon
       # is uneconomical. So we only ever add.
       indices.each do |index|
         progress.title = "Importing #{index} "
-        batch_size     = 1_000
+        batch_size     = options[:batch_size]
         slice_size     = (batch_size / options[:concurrency]).ceil
 
         index.adapter.default_scope.reorder(nil).find_in_batches(batch_size: batch_size) do |batch|