Plain is headed towards 1.0! Subscribe for development updates →

 1"""Type guards for database backend narrowing.
 2
 3These type guards allow type checkers to narrow BaseDatabaseWrapper types to
 4vendor-specific implementations based on runtime checks.
 5"""
 6
 7from __future__ import annotations
 8
 9from typing import TYPE_CHECKING, TypeIs
10
11if TYPE_CHECKING:
12    from plain.models.backends.base.base import BaseDatabaseWrapper
13    from plain.models.backends.mysql.base import MySQLDatabaseWrapper
14    from plain.models.backends.postgresql.base import PostgreSQLDatabaseWrapper
15    from plain.models.backends.sqlite3.base import SQLiteDatabaseWrapper
16
17
18def is_sqlite_connection(
19    connection: BaseDatabaseWrapper,
20) -> TypeIs[SQLiteDatabaseWrapper]:
21    """Type guard to narrow BaseDatabaseWrapper to SQLiteDatabaseWrapper.
22
23    Args:
24        connection: A database connection instance.
25
26    Returns:
27        True if the connection is a SQLite connection.
28
29    Example:
30        >>> if is_sqlite_connection(connection):
31        >>>     # connection is now SQLiteDatabaseWrapper
32        >>>     connection.ops.jsonfield_datatype_values
33    """
34    return connection.vendor == "sqlite"
35
36
37def is_mysql_connection(
38    connection: BaseDatabaseWrapper,
39) -> TypeIs[MySQLDatabaseWrapper]:
40    """Type guard to narrow BaseDatabaseWrapper to MySQLDatabaseWrapper.
41
42    Args:
43        connection: A database connection instance.
44
45    Returns:
46        True if the connection is a MySQL/MariaDB connection.
47    """
48    return connection.vendor == "mysql"
49
50
51def is_postgresql_connection(
52    connection: BaseDatabaseWrapper,
53) -> TypeIs[PostgreSQLDatabaseWrapper]:
54    """Type guard to narrow BaseDatabaseWrapper to PostgreSQLDatabaseWrapper.
55
56    Args:
57        connection: A database connection instance.
58
59    Returns:
60        True if the connection is a PostgreSQL connection.
61    """
62    return connection.vendor == "postgresql"