From 214b0c7439d67b423815515c49eda8ecf0d62bfa Mon Sep 17 00:00:00 2001
From: Anders Kaseorg <andersk@mit.edu>
Date: Tue, 28 Jun 2022 13:44:20 -0700
Subject: [PATCH] Support cursor.execute(psycopg2.sql.Composable) (#1029)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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: https://github.com/python/typeshed/pull/7494

Signed-off-by: Anders Kaseorg <andersk@mit.edu>
---
 django-stubs/db/backends/postgresql/base.pyi |  3 ++-
 django-stubs/db/backends/utils.pyi           | 28 +++++++++++++++++---
 tests/typecheck/db/test_connection.yml       | 10 +++++++
 3 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/django-stubs/db/backends/postgresql/base.pyi b/django-stubs/db/backends/postgresql/base.pyi
index eb2722042..a59100e00 100644
--- a/django-stubs/db/backends/postgresql/base.pyi
+++ b/django-stubs/db/backends/postgresql/base.pyi
@@ -3,6 +3,7 @@ from typing import Any, Dict, Tuple, Type
 
 from django.db.backends.base.base import BaseDatabaseWrapper
 from django.db.backends.utils import CursorDebugWrapper as BaseCursorDebugWrapper
+from django.db.backends.utils import _ExecuteQuery
 
 from .client import DatabaseClient
 from .creation import DatabaseCreation
@@ -37,5 +38,5 @@ class DatabaseWrapper(BaseDatabaseWrapper):
     def pg_version(self) -> int: ...
 
 class CursorDebugWrapper(BaseCursorDebugWrapper):
-    def copy_expert(self, sql: str, file: IOBase, *args: Any): ...
+    def copy_expert(self, sql: _ExecuteQuery, file: IOBase, *args: Any): ...
     def copy_to(self, file: IOBase, table: str, *args: Any, **kwargs: Any): ...
diff --git a/django-stubs/db/backends/utils.pyi b/django-stubs/db/backends/utils.pyi
index fc5b7946e..be872ad2a 100644
--- a/django-stubs/db/backends/utils.pyi
+++ b/django-stubs/db/backends/utils.pyi
@@ -4,7 +4,21 @@ import types
 from contextlib import contextmanager
 from decimal import Decimal
 from logging import Logger
-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):
@@ -14,6 +28,14 @@ else:
 
 logger: Logger
 
+# 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 = Union[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]
@@ -37,8 +59,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
diff --git a/tests/typecheck/db/test_connection.yml b/tests/typecheck/db/test_connection.yml
index 6b3f88bf1..8130085d2 100644
--- a/tests/typecheck/db/test_connection.yml
+++ b/tests/typecheck/db/test_connection.yml
@@ -4,6 +4,16 @@
       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:
+        cursor.execute(SQL("INSERT INTO {} VALUES (%s)").format(Identifier("my_table")), [123])
+
+
 -   case: raw_connections
     main: |
       from django.db import connections