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