Fix default port on secure connection + Query Parameterization

asynch/connection.py

@@ -15,38 +15,44 @@ def __init__(
         user: str = constants.DEFAULT_USER,
         password: str = constants.DEFAULT_PASSWORD,
         host: str = constants.DEFAULT_HOST,
-        port: int = constants.DEFAULT_PORT,
+        port: int | None = None,
         database: str = constants.DEFAULT_DATABASE,
         cursor_cls=Cursor,
         echo: bool = False,
         stack_track: bool = False,
+        secure: bool = False,
         **kwargs,
     ):
+        if port is None:
+            port = constants.DEFAULT_SECURE_PORT if secure else constants.DEFAULT_PORT
         if dsn:
             config = parse_dsn(dsn)
-            self._connection = ProtoConnection(**config, stack_track=stack_track, **kwargs)
-            user = config.get("user", None) or user
-            password = config.get("password", None) or password
-            host = config.get("host", None) or host
-            port = config.get("port", None) or port
-            database = config.get("database", None) or database
+            user = config.get("user") or user
+            password = config.get("password") or password
+            host = config.get("host") or host
+            port = config.get("port") or port
+            database = config.get("database") or database
+            secure = config.get("secure") or secure
         else:
-            self._connection = ProtoConnection(
-                host=host,
-                port=port,
-                database=database,
-                user=user,
-                password=password,
-                stack_track=stack_track,
-                **kwargs,
-            )
+            config = {
+                "host": host,
+                "port": port,
+                "database": database,
+                "user": user,
+                "password": password,
+                "secure": secure,
+                "stack_track": stack_track,
+            }
+
+        self._connection = ProtoConnection(**config, stack_track=stack_track, **kwargs)
         self._dsn = dsn
         # dsn parts
         self._user = user
         self._password = password
         self._host = host
         self._port = port
         self._database = database
+        self._secure = secure
         # connection additional settings
         self._opened: bool = False
         self._closed: bool = False

asynch/proto/connection.py

@@ -161,10 +161,16 @@ def __init__(  # nosec:B107
                 )
             ),
             "strings_as_bytes": self.settings.pop("strings_as_bytes", False),
-            "strings_encoding": self.settings.pop("strings_encoding", constants.STRINGS_ENCODING),
+            "strings_encoding": self.settings.pop(
+                "strings_encoding", constants.STRINGS_ENCODING
+            ),
             "use_numpy": self.settings.pop("use_numpy", False),
-            "opentelemetry_traceparent": self.settings.pop("opentelemetry_traceparent", None),
-            "opentelemetry_tracestate": self.settings.pop("opentelemetry_tracestate", ""),
+            "opentelemetry_traceparent": self.settings.pop(
+                "opentelemetry_traceparent", None
+            ),
+            "opentelemetry_tracestate": self.settings.pop(
+                "opentelemetry_tracestate", ""
+            ),
             "quota_key": self.settings.pop("quota_key", ""),
             "input_format_null_as_default": self.settings.pop(
                 "input_format_null_as_default", False
@@ -459,7 +465,9 @@ async def _receive_packet(self):
         elif packet_type == ServerPacket.END_OF_STREAM:
             self.is_query_executing = False
         elif packet_type == ServerPacket.TABLE_COLUMNS:
-            packet.multistring_message = await self.receive_multistring_message(packet_type)
+            packet.multistring_message = await self.receive_multistring_message(
+                packet_type
+            )
         elif packet_type == ServerPacket.PART_UUIDS:
             packet.block = await self.receive_data()
 
@@ -487,17 +495,25 @@ async def packet_generator(self) -> AsyncGenerator:
 
             yield packet
 
-    async def receive_result(self, with_column_types=False, progress=False, columnar=False):
+    async def receive_result(
+        self, with_column_types=False, progress=False, columnar=False
+    ):
         generator = self.packet_generator()
 
         if progress:
             return ProgressQueryResult(
-                self.reader, generator, with_column_types=with_column_types, columnar=columnar
+                self.reader,
+                generator,
+                with_column_types=with_column_types,
+                columnar=columnar,
             )
 
         else:
             result = QueryResult(
-                self.reader, generator, with_column_types=with_column_types, columnar=columnar
+                self.reader,
+                generator,
+                with_column_types=with_column_types,
+                columnar=columnar,
             )
             return await result.get_result()
 
@@ -525,7 +541,10 @@ async def send_query(self, query: str, query_id: str = ""):
             revision >= constants.DBMS_MIN_REVISION_WITH_SETTINGS_SERIALIZED_AS_STRINGS
         )
         await write_settings(
-            self.writer, self.context.settings, settings_as_strings, self.settings_is_important
+            self.writer,
+            self.context.settings,
+            settings_as_strings,
+            self.settings_is_important,
         )
         if revision >= constants.DBMS_MIN_REVISION_WITH_INTERSERVER_SECRET:
             await self.writer.write_str("")
@@ -540,7 +559,9 @@ async def send_query(self, query: str, query_id: str = ""):
 
     async def _init_connection(self, host: str, port: int):
         self.host, self.port = host, port
-        reader, writer = await asyncio.open_connection(host, port, ssl=self._get_ssl_context())
+        reader, writer = await asyncio.open_connection(
+            host, port, ssl=self._get_ssl_context()
+        )
         self.writer = BufferedWriter(writer)
         self.reader = BufferedReader(reader)
         self.block_reader = self.get_block_reader()
@@ -584,14 +605,14 @@ async def connect(self):
 
     async def execute(
         self,
-        query,
-        args=None,
-        with_column_types=False,
-        external_tables=None,
-        query_id="",
-        settings=None,
-        types_check=False,
-        columnar=False,
+        query: str,
+        args: Iterable | None = None,
+        with_column_types: bool = False,
+        external_tables: Iterable[Mapping] | None = None,
+        query_id: str = "",
+        settings: dict | None = None,
+        types_check: bool = False,
+        columnar: bool = False,
     ):
         """
         Executes query.
@@ -721,7 +742,7 @@ async def process_ordinary_query_with_progress(
         columnar: bool = False,
     ):
         if params is not None:
-            query = self.substitute_params(query, params)
+            query = substitute_params(query, params)
 
         await self.send_query(query, query_id=query_id)
         await self.send_external_tables(external_tables, types_check=types_check)
@@ -773,18 +794,22 @@ async def process_ordinary_query(
         columnar: bool = False,
     ):
         if params is not None:
-            query = self.substitute_params(query, params)
+            query = substitute_params(query, params)
 
         await self.send_query(query, query_id=query_id)
         await self.send_external_tables(external_tables, types_check=types_check)
-        return await self.receive_result(with_column_types=with_column_types, columnar=columnar)
+        return await self.receive_result(
+            with_column_types=with_column_types, columnar=columnar
+        )
 
     async def send_external_tables(self, tables, types_check=False):
         for table in tables or []:
             if not table["structure"]:
                 raise ValueError(f'Empty table "{table["name"]}" structure')
 
-            block = RowOrientedBlock(table["structure"], table["data"], types_check=types_check)
+            block = RowOrientedBlock(
+                table["structure"], table["data"], types_check=types_check
+            )
             await self.send_block(block, table_name=table["name"])
 
         # Empty block, end of data transfer.
@@ -880,7 +905,7 @@ async def iter_process_ordinary_query(
         types_check: bool = False,
     ):
         if params is not None:
-            query = self.substitute_params(query, params)
+            query = substitute_params(query, params)
 
         await self.send_query(query, query_id=query_id)
         await self.send_external_tables(external_tables, types_check=types_check)

tests/test_proto/utils/test_escape.py

@@ -0,0 +1,90 @@
+import ssl
+from contextlib import nullcontext as does_not_raise
+from datetime import date, datetime
+from enum import Enum
+from typing import Any, ContextManager, Optional
+from uuid import UUID
+
+import pytest
+from pytest import param
+
+from asynch.proto.utils.escape import escape_param, escape_params, substitute_params
+
+
+@pytest.mark.parametrize(
+    ("item", "expected"),
+    [
+        param(None, "NULL", id="None"),
+        param(date(2025, 5, 21), "'2025-05-21'", id="date"),
+        param(datetime(2025, 5, 21, 12, 0, 0), "'2025-05-21 12:00:00'", id="datetime"),
+        param("test", "'test'", id="str"),
+        param([1, 2, 3], "[1, 2, 3]", id="list"),
+        param((1, 2, 3), "(1, 2, 3)", id="tuple"),
+        param(
+            UUID("123e4567-e89b-12d3-a456-426614174000"),
+            "'123e4567-e89b-12d3-a456-426614174000'",
+            id="uuid",
+        ),
+        param(Enum("testEnum", {"test1": "test1", "test2": "test2"}).test1, "'test1'", id="enum"),
+    ],
+)
+def test_escape_param(item, expected):
+    assert escape_param(item) == expected
+
+
+@pytest.mark.parametrize(
+    ("params", "expected"),
+    [
+        param({"test": "test"}, {"test": "'test'"}, id="strings"),
+        param(
+            {"test": 1, "test2": 2, "test3": 0, "test4": -1},
+            {"test": 1, "test2": 2, "test3": 0, "test4": -1},
+            id="ints",
+        ),
+        param(
+            {"test": 1.0, "test2": 0.1, "test3": -0.1, "test4": 0.0},
+            {"test": 1.0, "test2": 0.1, "test3": -0.1, "test4": 0.0},
+            id="floats",
+        ),
+        param({"test": None}, {"test": "NULL"}, id="none"),
+        param({"test": date(2025, 5, 21)}, {"test": "'2025-05-21'"}, id="dates"),
+        param(
+            {"test": datetime(2025, 5, 21, 12, 0, 0)},
+            {"test": "'2025-05-21 12:00:00'"},
+            id="datetimes",
+        ),
+        param(
+            {"test": UUID("123e4567-e89b-12d3-a456-426614174000")},
+            {"test": "'123e4567-e89b-12d3-a456-426614174000'"},
+            id="uuids",
+        ),
+    ],
+)
+def test_escape_params(params, expected):
+    assert escape_params(params) == expected
+
+
+@pytest.mark.parametrize(
+    ("query", "params", "expected"),
+    [
+        param("{test}", {"test": "test"}, "'test'", id="string"),
+        param("{test} {test}", {"test": "test"}, "'test' 'test'", id="strings"),
+        param("{test} {test}", {"test": None}, "NULL NULL", id="none"),
+        param("{test} {test}", {"test": 1}, "1 1", id="int"),
+        param("{test} {test}", {"test": 1.0}, "1.0 1.0", id="float"),
+        param("{test} {test}", {"test": date(2025, 5, 21)}, "'2025-05-21' '2025-05-21'", id="date"),
+        param(
+            "{test} {test}",
+            {"test": datetime(2025, 5, 21, 12, 0, 0)},
+            "'2025-05-21 12:00:00' '2025-05-21 12:00:00'",
+            id="datetime",
+        ),
+        param(
+            "{test} {test}",
+            {"test": UUID("123e4567-e89b-12d3-a456-426614174000")},
+            "'123e4567-e89b-12d3-a456-426614174000' '123e4567-e89b-12d3-a456-426614174000'",
+        ),
+    ],
+)
+def test_substitute_params(query, params, expected):
+    assert substitute_params(query, params) == expected