1import sys
 2
 3from psycopg import errors
 4
 5from plain.exceptions import ImproperlyConfigured
 6from plain.models.backends.base.creation import BaseDatabaseCreation
 7
 8
 9class DatabaseCreation(BaseDatabaseCreation):
10    def _quote_name(self, name):
11        return self.connection.ops.quote_name(name)
12
13    def _get_database_create_suffix(self, encoding=None, template=None):
14        suffix = ""
15        if encoding:
16            suffix += f" ENCODING '{encoding}'"
17        if template:
18            suffix += f" TEMPLATE {self._quote_name(template)}"
19        return suffix and "WITH" + suffix
20
21    def sql_table_creation_suffix(self):
22        test_settings = self.connection.settings_dict["TEST"]
23        if test_settings.get("COLLATION") is not None:
24            raise ImproperlyConfigured(
25                "PostgreSQL does not support collation setting at database "
26                "creation time."
27            )
28        return self._get_database_create_suffix(
29            encoding=test_settings["CHARSET"],
30            template=test_settings.get("TEMPLATE"),
31        )
32
33    def _execute_create_test_db(self, cursor, parameters):
34        try:
35            super()._execute_create_test_db(cursor, parameters)
36        except Exception as e:
37            cause = e.__cause__
38            if cause and not isinstance(cause, errors.DuplicateDatabase):
39                # All errors except "database already exists" cancel tests.
40                self.log(f"Got an error creating the test database: {e}")
41                sys.exit(2)
42            else:
43                raise