1  import os
  2  import re
  3  import argparse
  4  import tempfile
  5  import tarfile
  6  import zipfile
  7  from functools import wraps
  8  
  9  import gitlab
 10  
 11  
 12  class Gitlab(object):
 13      JOB_NAME_PATTERN = re.compile(r"(\w+)(\s+(\d+)/(\d+))?")
 14  
 15      DOWNLOAD_ERROR_MAX_RETRIES = 3
 16  
 17      def __init__(self, project_id=None):
 18          config_data_from_env = os.getenv("PYTHON_GITLAB_CONFIG")
 19          if config_data_from_env:
 20              # prefer to load config from env variable
 21              with tempfile.NamedTemporaryFile("w", delete=False) as temp_file:
 22                  temp_file.write(config_data_from_env)
 23              config_files = [temp_file.name]
 24          else:
 25              # otherwise try to use config file at local filesystem
 26              config_files = None
 27          gitlab_id = os.getenv("LOCAL_GITLAB_HTTPS_HOST")  # if None, will use the default gitlab server
 28          self.gitlab_inst = gitlab.Gitlab.from_config(gitlab_id=gitlab_id, config_files=config_files)
 29          self.gitlab_inst.auth()
 30          if project_id:
 31              self.project = self.gitlab_inst.projects.get(project_id)
 32          else:
 33              self.project = None
 34  
 35      def get_project_id(self, name, namespace=None):
 36          """
 37          search project ID by name
 38  
 39          :param name: project name
 40          :param namespace: namespace to match when we have multiple project with same name
 41          :return: project ID
 42          """
 43          projects = self.gitlab_inst.projects.list(search=name)
 44          for project in projects:
 45              if namespace is None:
 46                  if len(projects) == 1:
 47                      project_id = project.id
 48                      break
 49              if project.namespace["path"] == namespace:
 50                  project_id = project.id
 51                  break
 52          else:
 53              raise ValueError("Can't find project")
 54          return project_id
 55  
 56      def download_artifacts(self, job_id, destination):
 57          """
 58          download full job artifacts and extract to destination.
 59  
 60          :param job_id: Gitlab CI job ID
 61          :param destination: extract artifacts to path.
 62          """
 63          job = self.project.jobs.get(job_id)
 64  
 65          with tempfile.NamedTemporaryFile(delete=False) as temp_file:
 66              job.artifacts(streamed=True, action=temp_file.write)
 67  
 68          with zipfile.ZipFile(temp_file.name, "r") as archive_file:
 69              archive_file.extractall(destination)
 70  
 71      def retry_download(func):
 72          """
 73          This wrapper will only catch IOError and retry the whole function.
 74  
 75          So only use it with download functions, read() inside and atomic
 76          functions
 77          """
 78          @wraps(func)
 79          def wrapper(self, *args, **kwargs):
 80              retried = 0
 81              while True:
 82                  try:
 83                      res = func(self, *args, **kwargs)
 84                  except (IOError, EOFError) as e:
 85                      retried += 1
 86                      if retried > self.DOWNLOAD_ERROR_MAX_RETRIES:
 87                          raise e  # get out of the loop
 88                      else:
 89                          print('Retried for the {} time'.format(retried))
 90                          continue
 91                  else:
 92                      break
 93              return res
 94          return wrapper
 95  
 96      def download_artifact(self, job_id, artifact_path, destination=None):
 97          """
 98          download specific path of job artifacts and extract to destination.
 99  
100          :param job_id: Gitlab CI job ID
101          :param artifact_path: list of path in artifacts (relative path to artifact root path)
102          :param destination: destination of artifact. Do not save to file if destination is None
103          :return: A list of artifact file raw data.
104          """
105          job = self.project.jobs.get(job_id)
106  
107          raw_data_list = []
108  
109          for a_path in artifact_path:
110              try:
111                  data = job.artifact(a_path)
112              except gitlab.GitlabGetError as e:
113                  print("Failed to download '{}' from job {}".format(a_path, job_id))
114                  raise e
115              raw_data_list.append(data)
116              if destination:
117                  file_path = os.path.join(destination, a_path)
118                  try:
119                      os.makedirs(os.path.dirname(file_path))
120                  except OSError:
121                      # already exists
122                      pass
123                  with open(file_path, "wb") as f:
124                      f.write(data)
125  
126          return raw_data_list
127  
128      def find_job_id(self, job_name, pipeline_id=None, job_status="success"):
129          """
130          Get Job ID from job name of specific pipeline
131  
132          :param job_name: job name
133          :param pipeline_id: If None, will get pipeline id from CI pre-defined variable.
134          :param job_status: status of job. One pipeline could have multiple jobs with same name after retry.
135                             job_status is used to filter these jobs.
136          :return: a list of job IDs (parallel job will generate multiple jobs)
137          """
138          job_id_list = []
139          if pipeline_id is None:
140              pipeline_id = os.getenv("CI_PIPELINE_ID")
141          pipeline = self.project.pipelines.get(pipeline_id)
142          jobs = pipeline.jobs.list(all=True)
143          for job in jobs:
144              match = self.JOB_NAME_PATTERN.match(job.name)
145              if match:
146                  if match.group(1) == job_name and job.status == job_status:
147                      job_id_list.append({"id": job.id, "parallel_num": match.group(3)})
148          return job_id_list
149  
150      @retry_download
151      def download_archive(self, ref, destination, project_id=None):
152          """
153          Download archive of certain commit of a repository and extract to destination path
154  
155          :param ref: commit or branch name
156          :param destination: destination path of extracted archive file
157          :param project_id: download project of current instance if project_id is None
158          :return: root path name of archive file
159          """
160          if project_id is None:
161              project = self.project
162          else:
163              project = self.gitlab_inst.projects.get(project_id)
164  
165          with tempfile.NamedTemporaryFile(delete=False) as temp_file:
166              try:
167                  project.repository_archive(sha=ref, streamed=True, action=temp_file.write)
168              except gitlab.GitlabGetError as e:
169                  print("Failed to archive from project {}".format(project_id))
170                  raise e
171  
172          print("archive size: {:.03f}MB".format(float(os.path.getsize(temp_file.name)) / (1024 * 1024)))
173  
174          with tarfile.open(temp_file.name, "r") as archive_file:
175              root_name = archive_file.getnames()[0]
176              archive_file.extractall(destination)
177  
178          return os.path.join(os.path.realpath(destination), root_name)
179  
180  
181  if __name__ == '__main__':
182      parser = argparse.ArgumentParser()
183      parser.add_argument("action")
184      parser.add_argument("project_id", type=int)
185      parser.add_argument("--pipeline_id", "-i", type=int, default=None)
186      parser.add_argument("--ref", "-r", default="master")
187      parser.add_argument("--job_id", "-j", type=int, default=None)
188      parser.add_argument("--job_name", "-n", default=None)
189      parser.add_argument("--project_name", "-m", default=None)
190      parser.add_argument("--destination", "-d", default=None)
191      parser.add_argument("--artifact_path", "-a", nargs="*", default=None)
192      args = parser.parse_args()
193  
194      gitlab_inst = Gitlab(args.project_id)
195      if args.action == "download_artifacts":
196          gitlab_inst.download_artifacts(args.job_id, args.destination)
197      if args.action == "download_artifact":
198          gitlab_inst.download_artifact(args.job_id, args.artifact_path, args.destination)
199      elif args.action == "find_job_id":
200          job_ids = gitlab_inst.find_job_id(args.job_name, args.pipeline_id)
201          print(";".join([",".join([str(j["id"]), j["parallel_num"]]) for j in job_ids]))
202      elif args.action == "download_archive":
203          gitlab_inst.download_archive(args.ref, args.destination)
204      elif args.action == "get_project_id":
205          ret = gitlab_inst.get_project_id(args.project_name)
206          print("project id: {}".format(ret))