From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Fri, 27 Mar 2020 06:55:43 +0000 (+1300)
Subject: Allow b and s formats on query split
X-Git-Tag: 3.0.dev0~669
X-Git-Url: http://git.ipfire.org/gitweb.cgi?a=commitdiff_plain;h=1ca5e9e833907c8964bda1b2dc9746574570e3a9;p=thirdparty%2Fpsycopg.git

Allow b and s formats on query split
---

diff --git a/psycopg3/utils/queries.py b/psycopg3/utils/queries.py
index 6088fb441..f692294e0 100644
--- a/psycopg3/utils/queries.py
+++ b/psycopg3/utils/queries.py
@@ -16,8 +16,10 @@ def query2pg(query, vars, codec):
 
     - Convert Python placeholders (``%s``, ``%(name)s``) into Postgres
       format (``$1``, ``$2``)
-    - return ``query`` (bytes), ``order`` (sequence of names used in the
-      query, in the position they appear, in case of named params, else None)
+    - placeholders can be %s or %b (text or binary)
+    - return ``query`` (bytes), ``formats`` (list of formats) ``order``
+      (sequence of names used in the query, in the position they appear, in
+
     """
     if not isinstance(query, bytes):
         # encoding from str already happened
@@ -91,24 +93,25 @@ def split_query(query, encoding="ascii"):
     m = None
     for m in _re_placeholder.finditer(query):
         pre = query[cur : m.span(0)[0]]
-        parts.append([pre, m])
+        parts.append([pre, m, None])
         cur = m.span(0)[1]
     if m is None:
-        parts.append([query, None])
+        parts.append([query, None, None])
     else:
-        parts.append([query[cur:], None])
+        parts.append([query[cur:], None, None])
 
     # drop the "%%", validate
     i = 0
     phtype = None
     while i < len(parts):
-        m = parts[i][1]
+        part = parts[i]
+        m = part[1]
         if m is None:
             break  # last part
         ph = m.group(0)
         if ph == b"%%":
             # unescape '%%' to '%' and merge the parts
-            parts[i + 1][0] = parts[i][0] + b"%" + parts[i + 1][0]
+            parts[i + 1][0] = part[0] + b"%" + parts[i + 1][0]
             del parts[i]
             continue
         if ph == b"%(":
@@ -122,22 +125,25 @@ def split_query(query, encoding="ascii"):
                 "incomplete placeholder: '%'; if you want to use '%' as an"
                 " operator you can double it up, i.e. use '%%'"
             )
-        elif ph[-1:] != b"s":
+        elif ph[-1:] not in b"bs":
             raise exc.ProgrammingError(
-                f"only '%s' and '%(name)s' placeholders allowed, got"
+                f"only '%s' and '%b' placeholders allowed, got"
                 f" {m.group(0).decode(encoding)}"
             )
 
         # Index or name
         if m.group(1) is None:
-            parts[i][1] = i
+            part[1] = i
         else:
-            parts[i][1] = m.group(1)
+            part[1] = m.group(1)
+
+        # Binary format
+        part[2] = ph[-1:] == b"b"
 
         if phtype is None:
-            phtype = type(parts[i][1])
+            phtype = type(part[1])
         else:
-            if phtype is not type(parts[i][1]):  # noqa
+            if phtype is not type(part[1]):  # noqa
                 raise exc.ProgrammingError(
                     "positional and named placeholders cannot be mixed"
                 )
diff --git a/tests/test_query.py b/tests/test_query.py
index a7994cc34..8ad6a342e 100644
--- a/tests/test_query.py
+++ b/tests/test_query.py
@@ -8,17 +8,34 @@ from psycopg3.utils.queries import split_query, query2pg, reorder_params
 @pytest.mark.parametrize(
     "input, want",
     [
-        (b"", [[b"", None]]),
-        (b"foo bar", [[b"foo bar", None]]),
-        (b"foo %% bar", [[b"foo % bar", None]]),
-        (b"%s", [[b"", 0], [b"", None]]),
-        (b"%s foo", [[b"", 0], [b" foo", None]]),
-        (b"foo %s", [[b"foo ", 0], [b"", None]]),
-        (b"foo %%%s bar", [[b"foo %", 0], [b" bar", None]]),
-        (b"foo %(name)s bar", [[b"foo ", b"name"], [b" bar", None]]),
+        (b"", [[b"", None, None]]),
+        (b"foo bar", [[b"foo bar", None, None]]),
+        (b"foo %% bar", [[b"foo % bar", None, None]]),
+        (b"%s", [[b"", 0, False], [b"", None, None]]),
+        (b"%s foo", [[b"", 0, False], [b" foo", None, None]]),
+        (b"%b foo", [[b"", 0, True], [b" foo", None, None]]),
+        (b"foo %s", [[b"foo ", 0, False], [b"", None, None]]),
+        (b"foo %%%s bar", [[b"foo %", 0, False], [b" bar", None, None]]),
         (
-            b"foo %s%s bar %s baz",
-            [[b"foo ", 0], [b"", 1], [b" bar ", 2], [b" baz", None]],
+            b"foo %(name)s bar",
+            [[b"foo ", b"name", False], [b" bar", None, None]],
+        ),
+        (
+            b"foo %(name)s %(name)b bar",
+            [
+                [b"foo ", b"name", False],
+                [b" ", b"name", True],
+                [b" bar", None, None],
+            ],
+        ),
+        (
+            b"foo %s%b bar %s baz",
+            [
+                [b"foo ", 0, False],
+                [b"", 1, True],
+                [b" bar ", 2, False],
+                [b" baz", None, None],
+            ],
         ),
     ],
 )