gh_db_load.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113
  1. #!/usr/bin/env python3
  2. #
  3. # Copyright (c) 2021 Project CHIP Authors
  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. """Fetch data from GitHub size artifacts."""
  18. import io
  19. import logging
  20. import sys
  21. import memdf.sizedb
  22. import memdf.util.config
  23. import memdf.util.markdown
  24. import memdf.util.sqlite
  25. from memdf import Config, ConfigDescription
  26. from memdf.util.github import Gh
  27. GITHUB_CONFIG: ConfigDescription = {
  28. Config.group_def('github'): {
  29. 'title': 'github options',
  30. },
  31. 'github.event': {
  32. 'help': 'Download only event type(s) (default ‘push’)',
  33. 'metavar': 'EVENT',
  34. 'default': [],
  35. 'argparse': {
  36. 'alias': ['--event']
  37. },
  38. },
  39. 'github.limit-artifacts': {
  40. 'help': 'Download no more than COUNT artifacts',
  41. 'metavar': 'COUNT',
  42. 'default': 0,
  43. 'argparse': {
  44. 'type': int,
  45. },
  46. },
  47. 'github.label': {
  48. 'help': 'Download artifacts for one label only',
  49. 'metavar': 'LABEL',
  50. 'default': '',
  51. },
  52. }
  53. def main(argv):
  54. status = 0
  55. try:
  56. sqlite_config = memdf.util.sqlite.CONFIG
  57. sqlite_config['database.file']['argparse']['required'] = True
  58. config = Config().init({
  59. **memdf.util.config.CONFIG,
  60. **memdf.util.github.CONFIG,
  61. **sqlite_config,
  62. **GITHUB_CONFIG,
  63. })
  64. config.argparse.add_argument('inputs', metavar='FILE', nargs='*')
  65. config.parse(argv)
  66. db = memdf.sizedb.SizeDatabase(config['database.file']).open()
  67. if gh := Gh(config):
  68. artifact_limit = config['github.limit-artifacts']
  69. artifacts_added = 0
  70. events = config['github.event']
  71. if not events:
  72. events = ['push']
  73. for a in gh.get_size_artifacts(label=config['github.label']):
  74. if events and a.event not in events:
  75. logging.debug('Skipping %s artifact %d', a.event, a.id)
  76. continue
  77. cur = db.execute('SELECT id FROM build WHERE artifact = ?',
  78. (a.id,))
  79. if cur.fetchone():
  80. logging.debug('Skipping known artifact %d', a.id)
  81. continue
  82. blob = gh.download_artifact(a.id)
  83. if blob:
  84. logging.info('Adding artifact %d %s %s %s %s',
  85. a.id, a.commit[:12], a.pr, a.event, a.group)
  86. db.add_sizes_from_zipfile(io.BytesIO(blob),
  87. {'artifact': a.id})
  88. db.commit()
  89. artifacts_added += 1
  90. if artifact_limit and artifact_limit <= artifacts_added:
  91. break
  92. for filename in config['args.inputs']:
  93. db.add_sizes_from_file(filename)
  94. db.commit()
  95. except Exception as exception:
  96. raise exception
  97. return status
  98. if __name__ == '__main__':
  99. sys.exit(main(sys.argv))