Account for sql.SQL and sql.Composed Objects

pgtrigger/runtime.py

@@ -11,19 +11,23 @@
 from typing import TYPE_CHECKING, List, Union
 
 from django.db import connections
+from typing_extensions import TypeAlias
 
 from pgtrigger import registry, utils
 
 if utils.psycopg_maj_version == 2:
     import psycopg2.extensions
 elif utils.psycopg_maj_version == 3:
     import psycopg.pq
+    import psycopg.sql as psycopg_sql
 else:
     raise AssertionError
 
 if TYPE_CHECKING:
     from pgtrigger import Timing
 
+_Query: TypeAlias = "str | bytes | psycopg_sql.SQL | psycopg_sql.Composed"
+_Connection: TypeAlias = "psycopg.Connection | psycopg2.extensions.connection"
 
 # All triggers currently being ignored
 _ignore = threading.local()
@@ -32,12 +36,24 @@
 _schema = threading.local()
 
 
-def _is_concurrent_statement(sql: str | bytes) -> bool:
+def _query_to_str(query: _Query, connection: _Connection) -> str:
+    psycopg_3 = utils.psycopg_maj_version == 3
+    if isinstance(query, str):
+        return query
+    elif isinstance(query, bytes):
+        return query.decode()
+    elif psycopg_3 and isinstance(query, (psycopg_sql.SQL, psycopg_sql.Composed)):
+        return query.as_string(connection)
+    else:
+        raise AssertionError
+
+
+def _is_concurrent_statement(sql: _Query, connection: _Connection) -> bool:
     """
     True if the sql statement is concurrent and cannot be ran in a transaction
     """
+    sql = _query_to_str(sql, connection)
     sql = sql.strip().lower() if sql else ""
-    sql = sql.decode() if isinstance(sql, bytes) else sql
     return sql.startswith("create") and "concurrently" in sql
 
 
@@ -72,7 +88,7 @@ def _can_inject_variable(cursor, sql):
     """
     return (
         not getattr(cursor, "name", None)
-        and not _is_concurrent_statement(sql)
+        and not _is_concurrent_statement(sql, cursor.connection)
         and not _is_transaction_errored(cursor)
     )
 
@@ -92,7 +108,7 @@ def _inject_pgtrigger_ignore(execute, sql, params, many, context):
     """
     if _can_inject_variable(context["cursor"], sql):
         serialized_ignore = "{" + ",".join(_ignore.value) + "}"
-        sql = sql.decode() if isinstance(sql, bytes) else sql
+        sql = _query_to_str(sql, context["cursor"])
         sql = f"SELECT set_config('pgtrigger.ignore', %s, true); {sql}"
         params = [serialized_ignore, *(params or ())]
 

pgtrigger/tests/test_runtime.py

@@ -6,6 +6,14 @@
 import pgtrigger.utils
 from pgtrigger.tests import models, utils
 
+if pgtrigger.utils.psycopg_maj_version == 3:
+    from psycopg.sql import SQL, Literal
+else:
+    from unittest.mock import MagicMock
+
+    SQL = MagicMock()
+    Literal = MagicMock()
+
 
 @pytest.mark.django_db
 def test_schema():
@@ -228,20 +236,26 @@ def test_custom_db_table_ignore():
         assert not models.CustomTableName.objects.exists()
 
 
-@pytest.mark.skipif(
-    pgtrigger.utils.psycopg_maj_version == 3, reason="Psycopg2 preserves entire query"
-)
 @pytest.mark.django_db
 @pytest.mark.parametrize(
-    "sql, params",
+    "sql, params, min_psycopg_version",
     [
-        ("select count(*) from auth_user where id = %s", (1,)),
-        ("select count(*) from auth_user", ()),
-        (b"select count(*) from auth_user where id = %s", (1,)),
-        (b"select count(*) from auth_user", ()),
+        ("select count(*) from auth_user where id = %s", (1,), 2),
+        ("select count(*) from auth_user", (), 2),
+        (b"select count(*) from auth_user where id = %s", (1,), 2),
+        (b"select count(*) from auth_user", (), 2),
+        (SQL("select count(*) from auth_user where id = %s"), (1,), 3),
+        (  # Formatting creates a composed object
+            SQL("select count(*) from auth_user where id = {id}").format(id=1),
+            (),
+            3,
+        ),
     ],
 )
-def test_inject_trigger_ignore(settings, mocker, sql, params):
+def test_inject_trigger_ignore(settings, mocker, sql, params, min_psycopg_version):
+    if pgtrigger.utils.psycopg_maj_version < min_psycopg_version:
+        pytest.skip("Psycopg version is less than {}".format(min_psycopg_version))
+
     settings.DEBUG = True
     expected_sql_base = "SELECT set_config('pgtrigger.ignore', '{ignored_triggers}', true)"
     # Order isn't deterministic, so we need to check for either order.
@@ -255,7 +269,6 @@ def test_inject_trigger_ignore(settings, mocker, sql, params):
         with connection.cursor() as cursor:
             cursor.execute(sql, params)
             query = connection.queries[-1]
-
             assert query["sql"].startswith(expected_sql_1) or query["sql"].startswith(
                 expected_sql_2
             )