Development of an internal social media platform with personalised dashboards for students
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

creation.py 3.3KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677
  1. import sys
  2. from psycopg2 import errorcodes
  3. from django.db.backends.base.creation import BaseDatabaseCreation
  4. from django.db.backends.utils import strip_quotes
  5. class DatabaseCreation(BaseDatabaseCreation):
  6. def _quote_name(self, name):
  7. return self.connection.ops.quote_name(name)
  8. def _get_database_create_suffix(self, encoding=None, template=None):
  9. suffix = ""
  10. if encoding:
  11. suffix += " ENCODING '{}'".format(encoding)
  12. if template:
  13. suffix += " TEMPLATE {}".format(self._quote_name(template))
  14. return suffix and "WITH" + suffix
  15. def sql_table_creation_suffix(self):
  16. test_settings = self.connection.settings_dict['TEST']
  17. assert test_settings['COLLATION'] is None, (
  18. "PostgreSQL does not support collation setting at database creation time."
  19. )
  20. return self._get_database_create_suffix(
  21. encoding=test_settings['CHARSET'],
  22. template=test_settings.get('TEMPLATE'),
  23. )
  24. def _database_exists(self, cursor, database_name):
  25. cursor.execute('SELECT 1 FROM pg_catalog.pg_database WHERE datname = %s', [strip_quotes(database_name)])
  26. return cursor.fetchone() is not None
  27. def _execute_create_test_db(self, cursor, parameters, keepdb=False):
  28. try:
  29. if keepdb and self._database_exists(cursor, parameters['dbname']):
  30. # If the database should be kept and it already exists, don't
  31. # try to create a new one.
  32. return
  33. super()._execute_create_test_db(cursor, parameters, keepdb)
  34. except Exception as e:
  35. if getattr(e.__cause__, 'pgcode', '') != errorcodes.DUPLICATE_DATABASE:
  36. # All errors except "database already exists" cancel tests.
  37. sys.stderr.write('Got an error creating the test database: %s\n' % e)
  38. sys.exit(2)
  39. elif not keepdb:
  40. # If the database should be kept, ignore "database already
  41. # exists".
  42. raise e
  43. def _clone_test_db(self, suffix, verbosity, keepdb=False):
  44. # CREATE DATABASE ... WITH TEMPLATE ... requires closing connections
  45. # to the template database.
  46. self.connection.close()
  47. source_database_name = self.connection.settings_dict['NAME']
  48. target_database_name = self.get_test_db_clone_settings(suffix)['NAME']
  49. test_db_params = {
  50. 'dbname': self._quote_name(target_database_name),
  51. 'suffix': self._get_database_create_suffix(template=source_database_name),
  52. }
  53. with self._nodb_connection.cursor() as cursor:
  54. try:
  55. self._execute_create_test_db(cursor, test_db_params, keepdb)
  56. except Exception as e:
  57. try:
  58. if verbosity >= 1:
  59. print("Destroying old test database for alias %s..." % (
  60. self._get_database_display_str(verbosity, target_database_name),
  61. ))
  62. cursor.execute('DROP DATABASE %(dbname)s' % test_db_params)
  63. self._execute_create_test_db(cursor, test_db_params, keepdb)
  64. except Exception as e:
  65. sys.stderr.write("Got an error cloning the test database: %s\n" % e)
  66. sys.exit(2)