database.py 4.4KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133
  1. # -*- coding: utf-8 -*-
  2. import argparse
  3. import plaster_pastedeploy
  4. import transaction
  5. from depot.manager import DepotManager
  6. from pyramid.paster import (
  7. get_appsettings,
  8. setup_logging,
  9. )
  10. from tracim import CFG
  11. from tracim.fixtures import FixturesLoader
  12. from tracim.fixtures.users_and_groups import Base as BaseFixture
  13. from tracim.fixtures.content import Content as ContentFixture
  14. from sqlalchemy.exc import IntegrityError
  15. from tracim.command import AppContextCommand
  16. from tracim.models.meta import DeclarativeBase
  17. from tracim.models import (
  18. get_engine,
  19. get_session_factory,
  20. get_tm_session,
  21. )
  22. class InitializeDBCommand(AppContextCommand):
  23. auto_setup_context = False
  24. def get_description(self) -> str:
  25. return "Initialize database"
  26. def get_parser(self, prog_name: str) -> argparse.ArgumentParser:
  27. parser = super().get_parser(prog_name)
  28. parser.add_argument(
  29. "--test-data",
  30. help='Add some default data to database to make test',
  31. dest='test_data',
  32. required=False,
  33. action='store_true',
  34. default=False,
  35. )
  36. return parser
  37. def take_action(self, parsed_args: argparse.Namespace) -> None:
  38. super(InitializeDBCommand, self).take_action(parsed_args)
  39. config_uri = parsed_args.config_file
  40. setup_logging(config_uri)
  41. settings = get_appsettings(config_uri)
  42. self._create_schema(settings)
  43. self._populate_database(settings, add_test_data=parsed_args.test_data)
  44. @classmethod
  45. def _create_schema(
  46. cls,
  47. settings: plaster_pastedeploy.ConfigDict
  48. ) -> None:
  49. print("- Create Schemas of databases -")
  50. engine = get_engine(settings)
  51. DeclarativeBase.metadata.create_all(engine)
  52. @classmethod
  53. def _populate_database(
  54. cls,
  55. settings: plaster_pastedeploy.ConfigDict,
  56. add_test_data: bool
  57. ) -> None:
  58. engine = get_engine(settings)
  59. session_factory = get_session_factory(engine)
  60. app_config = CFG(settings)
  61. print("- Populate database with default data -")
  62. with transaction.manager:
  63. dbsession = get_tm_session(session_factory, transaction.manager)
  64. try:
  65. fixtures = [BaseFixture]
  66. fixtures_loader = FixturesLoader(dbsession, app_config)
  67. fixtures_loader.loads(fixtures)
  68. transaction.commit()
  69. if add_test_data:
  70. app_config.configure_filedepot()
  71. fixtures = [ContentFixture]
  72. fixtures_loader.loads(fixtures)
  73. transaction.commit()
  74. print("Database initialized.")
  75. except IntegrityError:
  76. print('Warning, there was a problem when adding default data'
  77. ', it may have already been added:')
  78. import traceback
  79. print(traceback.format_exc())
  80. transaction.abort()
  81. print('Database initialization failed')
  82. class DeleteDBCommand(AppContextCommand):
  83. auto_setup_context = False
  84. def get_description(self) -> str:
  85. return "Delete database"
  86. def get_parser(self, prog_name: str) -> argparse.ArgumentParser:
  87. parser = super().get_parser(prog_name)
  88. parser.add_argument(
  89. "--force",
  90. help='force delete of database',
  91. dest='force',
  92. required=False,
  93. action='store_true',
  94. default=False,
  95. )
  96. return parser
  97. def take_action(self, parsed_args: argparse.Namespace) -> None:
  98. super(DeleteDBCommand, self).take_action(parsed_args)
  99. config_uri = parsed_args.config_file
  100. setup_logging(config_uri)
  101. settings = get_appsettings(config_uri)
  102. engine = get_engine(settings)
  103. app_config = CFG(settings)
  104. app_config.configure_filedepot()
  105. if parsed_args.force:
  106. print('Database deletion begin.')
  107. DeclarativeBase.metadata.drop_all(engine)
  108. print('Database deletion done.')
  109. print('Cleaning depot begin.')
  110. depot = DepotManager.get()
  111. depot_files = depot.list()
  112. for file_ in depot_files:
  113. depot.delete(file_)
  114. print('Cleaning depot done.')
  115. else:
  116. print('Warning, You should use --force if you really want to'
  117. ' delete database.')