Support cursor.execute(psycopg2.sql.Composable)

In addition to str, PostgreSQL cursors accept the
psycopg2.sql.Composable type, which is useful for guarding against SQL
injections when building raw queries that can’t be parameterized in
the normal way (e.g. interpolating identifiers).

In order to avoid reintroducing a dependency on psycopg2, we define a
Protocol that matches psycopg2.sql.Composable.

Documentation: https://www.psycopg.org/docs/sql.html
Related: python/typeshed#7494

Signed-off-by: Anders Kaseorg <[email protected]>

Author: andersk

django-stubs/db/backends/utils.pyi

@@ -3,7 +3,21 @@ import sys
 import types
 from contextlib import contextmanager
 from decimal import Decimal
-from typing import Any, Dict, Generator, Iterator, List, Mapping, Optional, Sequence, Tuple, Type, Union, overload
+from typing import (
+    Any,
+    Dict,
+    Generator,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    Protocol,
+    Sequence,
+    Tuple,
+    Type,
+    Union,
+    overload,
+)
 from uuid import UUID
 
 if sys.version_info < (3, 8):
@@ -13,6 +27,14 @@ else:
 
 logger: Any
 
+# Protocol matching psycopg2.sql.Composable, to avoid depending psycopg2
+class _Composable(Protocol):
+    def as_string(self, context: Any) -> str: ...
+    def __add__(self, other: _Composable) -> _Composable: ...
+    def __mul__(self, n: int) -> _Composable: ...
+
+_ExecuteQuery = str | _Composable
+
 # Python types that can be adapted to SQL.
 _SQLType = Union[
     None, bool, int, float, Decimal, str, bytes, datetime.date, datetime.datetime, UUID, Tuple[Any, ...], List[Any]
@@ -36,8 +58,8 @@ class CursorWrapper:
     def callproc(
         self, procname: str, params: Optional[Sequence[Any]] = ..., kparams: Optional[Dict[str, int]] = ...
     ) -> Any: ...
-    def execute(self, sql: str, params: _ExecuteParameters = ...) -> Any: ...
-    def executemany(self, sql: str, param_list: Sequence[_ExecuteParameters]) -> Any: ...
+    def execute(self, sql: _ExecuteQuery, params: _ExecuteParameters = ...) -> Any: ...
+    def executemany(self, sql: _ExecuteQuery, param_list: Sequence[_ExecuteParameters]) -> Any: ...
 
 class CursorDebugWrapper(CursorWrapper):
     cursor: Any

tests/typecheck/db/test_connection.yml

@@ -4,6 +4,13 @@
       with connection.cursor() as cursor:
         reveal_type(cursor)  # N: Revealed type is "django.db.backends.utils.CursorWrapper"
         cursor.execute("SELECT %s", [123])
+-   case: raw_connection_psycopg2_composable
+    main: |
+      from django.db import connection
+      from psycopg2.sql import SQL, Identifier
+      with connection.cursor() as cursor:
+        reveal_type(cursor)  # N: Revealed type is "django.db.backends.utils.CursorWrapper"
+        cursor.execute(SQL("INSERT INTO {} VALUES (%s)").format(Identifier("my_table")), [123])
 -   case: raw_connections
     main: |
       from django.db import connections