ImportTask.py 7.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224
  1. #!/usr/bin/env python3
  2. # Contest Management System - http://cms-dev.github.io/
  3. # Copyright © 2010-2013 Giovanni Mascellani <mascellani@poisson.phc.unipi.it>
  4. # Copyright © 2010-2018 Stefano Maggiolo <s.maggiolo@gmail.com>
  5. # Copyright © 2010-2012 Matteo Boscariol <boscarim@hotmail.com>
  6. # Copyright © 2013 Luca Wehrstedt <luca.wehrstedt@gmail.com>
  7. # Copyright © 2014-2016 William Di Luigi <williamdiluigi@gmail.com>
  8. # Copyright © 2015 Luca Chiodini <luca@chiodini.org>
  9. #
  10. # This program is free software: you can redistribute it and/or modify
  11. # it under the terms of the GNU Affero General Public License as
  12. # published by the Free Software Foundation, either version 3 of the
  13. # License, or (at your option) any later version.
  14. #
  15. # This program is distributed in the hope that it will be useful,
  16. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  17. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  18. # GNU Affero General Public License for more details.
  19. #
  20. # You should have received a copy of the GNU Affero General Public License
  21. # along with this program. If not, see <http://www.gnu.org/licenses/>.
  22. """This script imports a task from disk using one of the available
  23. loaders.
  24. The data parsed by the loader is used to create a new Task in the
  25. database.
  26. """
  27. # We enable monkey patching to make many libraries gevent-friendly
  28. # (for instance, urllib3, used by requests)
  29. import gevent.monkey
  30. gevent.monkey.patch_all() # noqa
  31. import argparse
  32. import logging
  33. import os
  34. import sys
  35. from cms import utf8_decoder
  36. from cms.db import SessionGen, Task
  37. from cms.db.filecacher import FileCacher
  38. from cmscontrib.importing import ImportDataError, contest_from_db, update_task
  39. from cmscontrib.loaders import choose_loader, build_epilog
  40. logger = logging.getLogger(__name__)
  41. class TaskImporter:
  42. """This script creates a task
  43. """
  44. def __init__(self, path, prefix, override_name, update, no_statement,
  45. contest_id, loader_class):
  46. """Create the importer object for a task.
  47. path (string): the path to the file or directory to import.
  48. prefix (string|None): an optional prefix added to the task name.
  49. override_name (string|None): an optional new name for the task.
  50. update (bool): if the task already exists, try to update it.
  51. no_statement (bool): do not try to import the task statement.
  52. contest_id (int|None): if set, the new task will be tied to this
  53. contest; if not set, the task will not be tied to any contest, or
  54. if this was an update, will remain tied to the previous contest.
  55. """
  56. self.file_cacher = FileCacher()
  57. self.prefix = prefix
  58. self.override_name = override_name
  59. self.update = update
  60. self.no_statement = no_statement
  61. self.contest_id = contest_id
  62. self.loader = loader_class(os.path.abspath(path), self.file_cacher)
  63. def do_import(self):
  64. """Get the task from the TaskLoader and store it."""
  65. # We need to check whether the task has changed *before* calling
  66. # get_task() as that method might reset the "has_changed" bit..
  67. task_has_changed = False
  68. if self.update:
  69. task_has_changed = self.loader.task_has_changed()
  70. # Get the task
  71. task = self.loader.get_task(get_statement=not self.no_statement)
  72. if task is None:
  73. return False
  74. # Override name, if necessary
  75. if self.override_name:
  76. task.name = self.override_name
  77. # Apply the prefix, if there is one
  78. if self.prefix:
  79. task.name = self.prefix + task.name
  80. # Store
  81. logger.info("Creating task on the database.")
  82. with SessionGen() as session:
  83. try:
  84. contest = contest_from_db(self.contest_id, session)
  85. task = self._task_to_db(
  86. session, contest, task, task_has_changed)
  87. except ImportDataError as e:
  88. logger.error(str(e))
  89. logger.info("Error while importing, no changes were made.")
  90. return False
  91. session.commit()
  92. task_id = task.id
  93. logger.info("Import finished (new task id: %s).", task_id)
  94. return True
  95. def _task_to_db(self, session, contest, new_task, task_has_changed):
  96. """Add the task to the DB
  97. Return the task, or raise in case of one of these errors:
  98. - if the task is not in the DB and user did not ask to update it;
  99. - if the task is already in the DB and attached to another contest.
  100. """
  101. task = session.query(Task).filter(Task.name == new_task.name).first()
  102. if task is None:
  103. if contest is not None:
  104. logger.info("Attaching task to contest (id %s.)",
  105. self.contest_id)
  106. new_task.num = len(contest.tasks)
  107. new_task.contest = contest
  108. session.add(new_task)
  109. return new_task
  110. if not self.update:
  111. raise ImportDataError(
  112. "Task \"%s\" already exists in database. "
  113. "Use --update to update it." % new_task.name)
  114. if contest is not None and task.contest_id != contest.id:
  115. raise ImportDataError(
  116. "Task \"%s\" already tied to another contest." % task.name)
  117. if task_has_changed:
  118. logger.info(
  119. "Task \"%s\" data has changed, updating it.", task.name)
  120. update_task(task, new_task, get_statements=not self.no_statement)
  121. else:
  122. logger.info("Task \"%s\" data has not changed.", task.name)
  123. return task
  124. def main():
  125. """Parse arguments and launch process."""
  126. parser = argparse.ArgumentParser(
  127. description="Import a new task or update an existing one in CMS.",
  128. epilog=build_epilog(),
  129. formatter_class=argparse.RawDescriptionHelpFormatter
  130. )
  131. parser.add_argument(
  132. "-L", "--loader",
  133. action="store", type=utf8_decoder,
  134. default=None,
  135. help="use the specified loader (default: autodetect)"
  136. )
  137. parser.add_argument(
  138. "-u", "--update",
  139. action="store_true",
  140. help="update an existing task"
  141. )
  142. parser.add_argument(
  143. "-S", "--no-statement",
  144. action="store_true",
  145. help="do not import / update task statement"
  146. )
  147. parser.add_argument(
  148. "-c", "--contest-id",
  149. action="store", type=int,
  150. help="id of the contest the task will be attached to"
  151. )
  152. parser.add_argument(
  153. "-p", "--prefix",
  154. action="store", type=utf8_decoder,
  155. help="the prefix to be added before the task name"
  156. )
  157. parser.add_argument(
  158. "-n", "--name",
  159. action="store", type=utf8_decoder,
  160. help="the new name that will override the task name"
  161. )
  162. parser.add_argument(
  163. "target",
  164. action="store", type=utf8_decoder,
  165. help="target file/directory from where to import task(s)"
  166. )
  167. args = parser.parse_args()
  168. loader_class = choose_loader(
  169. args.loader,
  170. args.target,
  171. parser.error
  172. )
  173. importer = TaskImporter(path=args.target,
  174. update=args.update,
  175. no_statement=args.no_statement,
  176. contest_id=args.contest_id,
  177. prefix=args.prefix,
  178. override_name=args.name,
  179. loader_class=loader_class)
  180. success = importer.do_import()
  181. return 0 if success is True else 1
  182. if __name__ == "__main__":
  183. sys.exit(main())