prepare_untracked_uploads.rb 4.49 KB
Newer Older
1
2
module Gitlab
  module BackgroundMigration
3
    class PrepareUntrackedUploads
4
5
6
      # For bulk_queue_background_migration_jobs_by_range
      include Database::MigrationHelpers

7
      FILE_PATH_BATCH_SIZE = 500
8
9
      RELATIVE_UPLOAD_DIR = "uploads".freeze
      ABSOLUTE_UPLOAD_DIR = "#{CarrierWave.root}/#{RELATIVE_UPLOAD_DIR}".freeze
Michael Kozono's avatar
Michael Kozono committed
10
      FOLLOW_UP_MIGRATION = 'PopulateUntrackedUploads'.freeze
11
12
13
      START_WITH_CARRIERWAVE_ROOT_REGEX = %r{\A#{CarrierWave.root}/}
      EXCLUDED_HASHED_UPLOADS_PATH = "#{ABSOLUTE_UPLOAD_DIR}/@hashed/*".freeze
      EXCLUDED_TMP_UPLOADS_PATH = "#{ABSOLUTE_UPLOAD_DIR}/tmp/*".freeze
14

15
      class UntrackedFile < ActiveRecord::Base
16
17
        include EachBatch

18
        self.table_name = 'untracked_files_for_uploads'
19
20
21
      end

      def perform
22
        ensure_temporary_tracking_table_exists
23
        store_untracked_file_paths
24
25
26
27
28
        schedule_populate_untracked_uploads_jobs
      end

      private

29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
      def ensure_temporary_tracking_table_exists
        unless UntrackedFile.connection.table_exists?(:untracked_files_for_uploads)
          UntrackedFile.connection.create_table :untracked_files_for_uploads do |t|
            t.string :path, limit: 600, null: false
            t.boolean :tracked, default: false, null: false
            t.timestamps_with_timezone null: false
          end
        end

        unless UntrackedFile.connection.index_exists?(:untracked_files_for_uploads, :path)
          UntrackedFile.connection.add_index :untracked_files_for_uploads, :path, unique: true
        end

        unless UntrackedFile.connection.index_exists?(:untracked_files_for_uploads, :tracked)
          UntrackedFile.connection.add_index :untracked_files_for_uploads, :tracked
        end
45
46
      end

47
      def store_untracked_file_paths
48
        return unless Dir.exist?(ABSOLUTE_UPLOAD_DIR)
49

50
        each_file_batch(ABSOLUTE_UPLOAD_DIR, FILE_PATH_BATCH_SIZE) do |file_paths|
51
          insert_file_paths(file_paths)
52
53
54
        end
      end

55
      def each_file_batch(search_dir, batch_size, &block)
56
        cmd = build_find_command(search_dir)
57

58
        Open3.popen2(*cmd) do |stdin, stdout, status_thread|
59
60
          yield_paths_in_batches(stdout, batch_size, &block)

61
62
63
64
          raise "Find command failed" unless status_thread.value.success?
        end
      end

65
66
67
68
      def yield_paths_in_batches(stdout, batch_size, &block)
        paths = []

        stdout.each_line("\0") do |line|
69
          paths << line.chomp("\0").sub(START_WITH_CARRIERWAVE_ROOT_REGEX, '')
70
71
72
73
74
75
76
77
78
79

          if paths.size >= batch_size
            yield(paths)
            paths = []
          end
        end

        yield(paths)
      end

80
      def build_find_command(search_dir)
81
        cmd = %W[find #{search_dir} -type f ! ( -path #{EXCLUDED_HASHED_UPLOADS_PATH} -prune ) ! ( -path #{EXCLUDED_TMP_UPLOADS_PATH} -prune ) -print0]
82

Michael Kozono's avatar
Michael Kozono committed
83
        cmd = %w[ionice -c Idle] + cmd if ionice_is_available?
84
85
86
87
88
89
90
91
92
93
94
95

        cmd
      end

      def ionice_is_available?
        Gitlab::Utils.which('ionice')
      rescue StandardError
        # In this case, returning false is relatively safe, even though it isn't very nice
        false
      end

      def insert_file_paths(file_paths)
Michael Kozono's avatar
Michael Kozono committed
96
97
98
99
        ActiveRecord::Base.transaction do
          file_paths.each do |file_path|
            insert_file_path(file_path)
          end
100
        end
101
102
      end

103
      def insert_file_path(file_path)
104
105
106
107
108
        if postgresql_pre_9_5?
          # No easy way to do ON CONFLICT DO NOTHING before Postgres 9.5 so just use Rails
          return UntrackedFile.where(path: file_path).first_or_create
        end

109
        table_columns_and_values = 'untracked_files_for_uploads (path, created_at, updated_at) VALUES (?, ?, ?)'
110

111
        sql = if postgresql?
Michael Kozono's avatar
Michael Kozono committed
112
113
114
115
                "INSERT INTO #{table_columns_and_values} ON CONFLICT DO NOTHING;"
              else
                "INSERT IGNORE INTO #{table_columns_and_values};"
              end
116
117

        timestamp = Time.now.utc.iso8601
Michael Kozono's avatar
Michael Kozono committed
118
        sql = ActiveRecord::Base.send(:sanitize_sql_array, [sql, file_path, timestamp, timestamp]) # rubocop:disable GitlabSecurity/PublicSend
119
120
121
        ActiveRecord::Base.connection.execute(sql)
      end

122
123
124
125
126
127
128
129
130
      def postgresql?
        @postgresql ||= Gitlab::Database.postgresql?
      end

      def postgresql_pre_9_5?
        @postgresql_pre_9_5 ||= postgresql? &&
          ActiveRecord::Base.connection.select_value('SHOW server_version_num').to_i < 90500
      end

131
      def schedule_populate_untracked_uploads_jobs
132
        bulk_queue_background_migration_jobs_by_range(UntrackedFile, FOLLOW_UP_MIGRATION)
133
134
135
136
      end
    end
  end
end