creation.py 3.8 KB

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