database.py 4.1KB

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