lib/github/wallet.nit

   1 # This file is part of NIT ( http://www.nitlanguage.org ).
   2 #
   3 # Copyright 2016 Alexandre Terrasa <alexandre@moz-code.org>
   4 #
   5 # Licensed under the Apache License, Version 2.0 (the "License");
   6 # you may not use this file except in compliance with the License.
   7 # You may obtain a copy of the License at
   8 #
   9 #     http://www.apache.org/licenses/LICENSE-2.0
  10 #
  11 # Unless required by applicable law or agreed to in writing, software
  12 # distributed under the License is distributed on an "AS IS" BASIS,
  13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14 # See the License for the specific language governing permissions and
  15 # limitations under the License.
  16
  17 # Github OAuth tokens management
  18 #
  19 # When using batch mode with the `github` API, we can rapidly reach the rate
  20 # limit allowed by Github.
  21 #
  22 # One solution consists in using a wallet of tokens so we can rely on more than
  23 # one token and switch them when one become exhausted.
  24 #
  25 # ## Using the Github wallet to check tokens
  26 #
  27 # One functionality of the wallet is to check the validity of a token against
  28 # the API. `check_token` will return false if a token is invalid or exhausted.
  29 #
  30 # ~~~
  31 # var wallet = new GithubWallet
  32 # assert not wallet.check_token("this is a bad token")
  33 # ~~~
  34 #
  35 # ## Storing tokens
  36 #
  37 # The wallet can also be used to store tokens and check all of them.
  38 #
  39 # ~~~
  40 # wallet = new GithubWallet
  41 # wallet.add "some token"
  42 # wallet.add "some other token"
  43 # ~~~
  44 #
  45 # or
  46 #
  47 # ~~~
  48 # wallet = new GithubWallet(["token 1", "token 2"])
  49 # ~~~
  50 #
  51 # The `show_status` method can be used to display a summary of the validity of
  52 # each token in the wallet.
  53 #
  54 # ~~~
  55 # wallet.show_status
  56 # ~~~
  57 #
  58 # Will display something like this:
  59 #
  60 # ~~~raw
  61 # Wallet (2 tokens):
  62 #  * [OK] token 1
  63 #  * [KO] token 2
  64 # ~~~
  65 #
  66 # ## Using the wallet to obtain a Github API client
  67 #
  68 # Using the wallet you can cycle through tokens and obtain a new Github API client
  69 # instance with a fresh rate limit.
  70 #
  71 # ~~~
  72 # wallet = new GithubWallet(["token 1", "token 2"])
  73 # var api = wallet.api
  74 # ~~~
  75 #
  76 # The wallet will automatically cycle through the registered tokens to find one
  77 # that works.
  78 #
  79 # If no valid token is found after all of them was tried, the wallet returns a
  80 # client based on the last tried one.
  81 module wallet
  82
  83 import github
  84 import logger
  85
  86 # Github OAuth tokens wallet
  87 class GithubWallet
  88
  89         # Github API tokens
  90         var tokens = new Array[String] is optional
  91
  92         # Logger used to display info about tokens state
  93         var logger = new Logger is optional, writable
  94
  95         # Add a new token in the wallet
  96         fun add(token: String) do tokens.add token
  97
  98         # Get an instance of GithubAPI based on the next available token.
  99         #
 100         # If no token is found, return an api based on the last exhausted token.
 101         fun api: GithubAPI do
 102                 var token
 103                 if tokens.is_empty then
 104                         logger.warn "No tokens, using `get_github_oauth`"
 105                         token = get_github_oauth
 106                 else
 107                         token = get_next_token
 108                         var tried = 0
 109                         while not check_token(token) do
 110                                 if tried >= tokens.length - 1 then
 111                                         logger.warn "Exhausted all tokens, using {token}"
 112                                         break
 113                                 end
 114                                 tried += 1
 115                                 token = get_next_token
 116                         end
 117                 end
 118                 var api = new GithubAPI(token)
 119                 api.enable_cache = true
 120                 return api
 121         end
 122
 123         # The current index in the `tokens` array
 124         private var current_index = 0
 125
 126         # The current token in the `tokens` array based on `current_index`
 127         fun current_token: String do return tokens[current_index]
 128
 129         # Get the next token in token `array` based on `current_token`.
 130         #
 131         # If the end of the list is reached, start again from the begining.
 132         fun get_next_token: String do
 133                 if tokens.is_empty then
 134                         return get_github_oauth
 135                 end
 136                 var token = current_token
 137
 138                 if current_index < tokens.length - 1 then
 139                         current_index += 1
 140                 else
 141                         current_index = 0
 142                 end
 143                 return token
 144         end
 145
 146         # Check if a token is valid
 147         fun check_token(token: String): Bool do
 148                 logger.debug "Try token {token}"
 149                 var api = new GithubAPI(token)
 150                 api.get_auth_user
 151                 return not api.was_error
 152         end
 153
 154         # Show wallet status in console
 155         fun show_status(no_color: nullable Bool) do
 156                 no_color = no_color or else false
 157
 158                 if tokens.is_empty then
 159                         print "Wallet is empty"
 160                         return
 161                 end
 162                 print "Wallet ({tokens.length} tokens):"
 163                 for token in tokens do
 164                         var status
 165                         if check_token(token) then
 166                                 status = if no_color then "OK" else "OK".green
 167                         else
 168                                 status = if no_color then "KO" else "KO".red
 169                         end
 170                         print " * [{status}] {token}"
 171                 end
 172         end
 173 end