Read about our upcoming Code of Conduct on this issue

client.rb 6.64 KB
Newer Older
1
2
# frozen_string_literal: true

Valery Sizov's avatar
Valery Sizov committed
3
module Gitlab
4
  module GithubImport
5
6
7
8
9
10
11
12
13
14
15
    # HTTP client for interacting with the GitHub API.
    #
    # This class is basically a fancy wrapped around Octokit while adding some
    # functionality to deal with rate limiting and parallel imports. Usage is
    # mostly the same as Octokit, for example:
    #
    #     client = GithubImport::Client.new('hunter2')
    #
    #     client.labels.each do |label|
    #       puts label.name
    #     end
Valery Sizov's avatar
Valery Sizov committed
16
    class Client
17
18
      include ::Gitlab::Utils::StrongMemoize

19
      attr_reader :octokit
20

21
22
      # A single page of data and the corresponding page number.
      Page = Struct.new(:objects, :number)
Valery Sizov's avatar
Valery Sizov committed
23

24
25
26
27
28
29
30
      # The minimum number of requests we want to keep available.
      #
      # We don't use a value of 0 as multiple threads may be using the same
      # token in parallel. This could result in all of them hitting the GitHub
      # rate limit at once. The threshold is put in place to not hit the limit
      # in most cases.
      RATE_LIMIT_THRESHOLD = 50
Valery Sizov's avatar
Valery Sizov committed
31

32
33
34
35
36
37
38
39
40
41
42
      # token - The GitHub API token to use.
      #
      # per_page - The number of objects that should be displayed per page.
      #
      # parallel - When set to true hitting the rate limit will result in a
      #            dedicated error being raised. When set to `false` we will
      #            instead just `sleep()` until the rate limit is reset. Setting
      #            this value to `true` for parallel importing is crucial as
      #            otherwise hitting the rate limit will result in a thread
      #            being blocked in a `sleep()` call for up to an hour.
      def initialize(token, per_page: 100, parallel: true)
43
44
45
46
47
48
49
50
        @octokit = Octokit::Client.new(
          access_token: token,
          per_page: per_page,
          api_endpoint: api_endpoint
        )

        @octokit.connection_options[:ssl] = { verify: verify_ssl }

51
        @parallel = parallel
52
53
      end

54
55
      def parallel?
        @parallel
56
      end
57

58
59
60
61
62
      # Returns the details of a GitHub user.
      #
      # username - The username of the user.
      def user(username)
        with_rate_limit { octokit.user(username) }
Valery Sizov's avatar
Valery Sizov committed
63
64
      end

65
66
67
68
69
      # Returns the details of a GitHub repository.
      #
      # name - The path (in the form `owner/repository`) of the repository.
      def repository(name)
        with_rate_limit { octokit.repo(name) }
Valery Sizov's avatar
Valery Sizov committed
70
71
      end

72
73
      def labels(*args)
        each_object(:labels, *args)
Valery Sizov's avatar
Valery Sizov committed
74
75
      end

76
77
      def milestones(*args)
        each_object(:milestones, *args)
Valery Sizov's avatar
Valery Sizov committed
78
79
      end

80
81
      def releases(*args)
        each_object(:releases, *args)
Valery Sizov's avatar
Valery Sizov committed
82
83
      end

84
85
86
87
88
89
90
91
92
      # Fetches data from the GitHub API and yields a Page object for every page
      # of data, without loading all of them into memory.
      #
      # method - The Octokit method to use for getting the data.
      # args - Arguments to pass to the Octokit method.
      #
      # rubocop: disable GitlabSecurity/PublicSend
      def each_page(method, *args, &block)
        return to_enum(__method__, method, *args) unless block_given?
93

94
95
96
97
98
99
        page =
          if args.last.is_a?(Hash) && args.last[:page]
            args.last[:page]
          else
            1
          end
100

101
102
        collection = with_rate_limit { octokit.public_send(method, *args) }
        next_url = octokit.last_response.rels[:next]
Valery Sizov's avatar
Valery Sizov committed
103

104
105
106
107
108
109
110
        yield Page.new(collection, page)

        while next_url
          response = with_rate_limit { next_url.get }
          next_url = response.rels[:next]

          yield Page.new(response.data, page += 1)
111
112
113
        end
      end

114
115
116
117
118
119
      # Iterates over all of the objects for the given method (e.g. `:labels`).
      #
      # method - The method to send to Octokit for querying data.
      # args - Any arguments to pass to the Octokit method.
      def each_object(method, *args, &block)
        return to_enum(__method__, method, *args) unless block_given?
Valery Sizov's avatar
Valery Sizov committed
120

121
122
123
124
        each_page(method, *args) do |page|
          page.objects.each do |object|
            yield object
          end
125
        end
Valery Sizov's avatar
Valery Sizov committed
126
      end
127

128
129
130
131
132
133
      # Yields the supplied block, responding to any rate limit errors.
      #
      # The exact strategy used for handling rate limiting errors depends on
      # whether we are running in parallel mode or not. For more information see
      # `#rate_or_wait_for_rate_limit`.
      def with_rate_limit
134
135
        return yield unless rate_limiting_enabled?

136
        request_count_counter.increment
137

138
        raise_or_wait_for_rate_limit unless requests_remaining?
139

140
141
142
143
        begin
          yield
        rescue Octokit::TooManyRequests
          raise_or_wait_for_rate_limit
144

145
146
147
148
          # This retry will only happen when running in sequential mode as we'll
          # raise an error in parallel mode.
          retry
        end
149
150
      end

151
152
153
      # Returns `true` if we're still allowed to perform API calls.
      def requests_remaining?
        remaining_requests > RATE_LIMIT_THRESHOLD
154
155
      end

156
157
158
      def remaining_requests
        octokit.rate_limit.remaining
      end
159

160
161
      def raise_or_wait_for_rate_limit
        rate_limit_counter.increment
162

163
164
        if parallel?
          raise RateLimitError
165
        else
166
          sleep(rate_limit_resets_in)
167
168
169
        end
      end

170
171
172
173
174
175
176
      def rate_limit_resets_in
        # We add a few seconds to the rate limit so we don't _immediately_
        # resume when the rate limit resets as this may result in us performing
        # a request before GitHub has a chance to reset the limit.
        octokit.rate_limit.resets_in + 5
      end

177
      def rate_limiting_enabled?
178
179
180
        strong_memoize(:rate_limiting_enabled) do
          api_endpoint.include?('.github.com')
        end
181
182
      end

183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
      def api_endpoint
        custom_api_endpoint || default_api_endpoint
      end

      def custom_api_endpoint
        github_omniauth_provider.dig('args', 'client_options', 'site')
      end

      def default_api_endpoint
        OmniAuth::Strategies::GitHub.default_options[:client_options][:site]
      end

      def verify_ssl
        github_omniauth_provider.fetch('verify_ssl', true)
      end

      def github_omniauth_provider
        @github_omniauth_provider ||=
          Gitlab.config.omniauth.providers
                .find { |provider| provider.name == 'github' }
                .to_h
204
205
206
207
208
209
210
211
212
213
214
215
216
217
      end

      def rate_limit_counter
        @rate_limit_counter ||= Gitlab::Metrics.counter(
          :github_importer_rate_limit_hits,
          'The number of times we hit the GitHub rate limit when importing projects'
        )
      end

      def request_count_counter
        @request_counter ||= Gitlab::Metrics.counter(
          :github_importer_request_count,
          'The number of GitHub API calls performed when importing projects'
        )
218
      end
Valery Sizov's avatar
Valery Sizov committed
219
220
221
    end
  end
end