Fix params substitution for select queries

asynch/proto/connection.py

@@ -4,7 +4,7 @@
 from collections.abc import AsyncGenerator
 from time import time
 from types import GeneratorType
-from typing import Optional, Union
+from typing import Any, Iterable, Mapping, Optional, Union
 from urllib.parse import urlparse
 
 from asynch.errors import (
@@ -712,13 +712,13 @@ async def execute_with_progress(
 
     async def process_ordinary_query_with_progress(
         self,
-        query,
-        params=None,
-        with_column_types=False,
-        external_tables=None,
-        query_id=None,
-        types_check=False,
-        columnar=False,
+        query: str,
+        params: Optional[Mapping[str, Any]] = None,
+        with_column_types: bool = False,
+        external_tables: Optional[Iterable[Mapping]] = None,
+        query_id: Optional[str] = None,
+        types_check: bool = False,
+        columnar: bool = False,
     ):
         if params is not None:
             query = self.substitute_params(query, params)
@@ -764,13 +764,13 @@ async def force_connect(self):
 
     async def process_ordinary_query(
         self,
-        query,
-        params=None,
-        with_column_types=False,
-        external_tables=None,
-        query_id="",
-        types_check=False,
-        columnar=False,
+        query: str,
+        params: Optional[Mapping[str, Any]] = None,
+        with_column_types: bool = False,
+        external_tables: Optional[Iterable[Mapping]] = None,
+        query_id: Optional[str] = None,
+        types_check: bool = False,
+        columnar: bool = False,
     ):
         if params is not None:
             query = self.substitute_params(query, params)
@@ -801,12 +801,15 @@ async def send_block(self, block, table_name=""):
 
         await self.block_writer.write(block)
 
-    def substitute_params(self, query, params):
-        if not isinstance(params, dict):
-            raise ValueError("Parameters are expected in dict form")
+    @staticmethod
+    def substitute_params(query: str, params: Mapping[str, Any]) -> str:
+        if not isinstance(params, Mapping):
+            raise ValueError("Parameters are expected to be a mapping")
 
-        escaped = escape_params(params)
-        return query % escaped
+        try:
+            return query.format(**escape_params(params))
+        except KeyError as exc:
+            raise KeyError(f"Parameter '{exc}' not found") from exc
 
     async def process_insert_query(
         self,
@@ -872,12 +875,12 @@ async def send_data(self, sample_block, data, types_check=False, columnar=False)
 
     async def iter_process_ordinary_query(
         self,
-        query,
-        params=None,
-        with_column_types=False,
-        external_tables=None,
-        query_id=None,
-        types_check=False,
+        query: str,
+        params: Optional[Mapping[str, Any]] = None,
+        with_column_types: bool = False,
+        external_tables: Optional[Iterable[Mapping]] = None,
+        query_id: Optional[str] = None,
+        types_check: bool = False,
     ):
         if params is not None:
             query = self.substitute_params(query, params)

asynch/proto/utils/escape.py

@@ -1,5 +1,6 @@
 from datetime import date, datetime
 from enum import Enum
+from typing import Any, Mapping
 from uuid import UUID
 
 from .compat import string_types, text_type
@@ -18,7 +19,7 @@
 }
 
 
-def escape_param(item):
+def escape_param(item: Any) -> str:
     if item is None:
         return "NULL"
 
@@ -44,13 +45,8 @@ def escape_param(item):
         return "'%s'" % str(item)
 
     else:
-        return item
+        return str(item)
 
 
-def escape_params(params):
-    escaped = {}
-
-    for key, value in params.items():
-        escaped[key] = escape_param(value)
-
-    return escaped
+def escape_params(params: Mapping[str, Any]) -> dict[str, str]:
+    return {key: escape_param(value) for key, value in params.items()}

tests/test_cursors.py

@@ -58,6 +58,10 @@ async def test_fetchone(conn: Connection):
         ret = await cursor.fetchone()
         assert ret == (1,)
 
+        await cursor.execute("SELECT {val}", args={"val": 2})
+        ret = await cursor.fetchone()
+        assert ret == (2,)
+
         await cursor.execute("SELECT * FROM system.tables")
         ret = await cursor.fetchall()
         assert isinstance(ret, list)
@@ -70,6 +74,10 @@ async def test_fetchall(conn: Connection):
         ret = await cursor.fetchall()
         assert ret == [(1,)]
 
+        await cursor.execute("SELECT {val}", args={"val": 2})
+        ret = await cursor.fetchall()
+        assert ret == [(2,)]
+
 
 @pytest.mark.asyncio
 async def test_dict_cursor(conn: Connection):
@@ -78,6 +86,10 @@ async def test_dict_cursor(conn: Connection):
         ret = await cursor.fetchall()
         assert ret == [{"1": 1}]
 
+        await cursor.execute("SELECT {val}", args={"val": 2})
+        ret = await cursor.fetchall()
+        assert ret == [{"2": 2}]
+
 
 @pytest.mark.asyncio
 async def test_insert_dict(conn: Connection):

tests/test_proto/test_proto_connection.py

@@ -89,6 +89,13 @@ async def test_execute(proto_conn: ProtoConnection):
     assert ret == [(1,)]
 
 
+@pytest.mark.asyncio
+async def test_execute_with_args(proto_conn: ProtoConnection):
+    query = "SELECT {val}"
+    ret = await proto_conn.execute(query, args={"val": 2})
+    assert ret == [(2,)]
+
+
 @asynccontextmanager
 async def create_table(connection, spec):
     await connection.execute("DROP TABLE IF EXISTS test.test")