diff --git a/asynch/connection.py b/asynch/connection.py index 27383dcc..a0b2493a 100644 --- a/asynch/connection.py +++ b/asynch/connection.py @@ -15,31 +15,36 @@ 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 @@ -47,6 +52,7 @@ def __init__( self._host = host self._port = port self._database = database + self._secure = secure # connection additional settings self._opened: bool = False self._closed: bool = False diff --git a/asynch/proto/connection.py b/asynch/proto/connection.py index e9b467bd..c6a1af5f 100644 --- a/asynch/proto/connection.py +++ b/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) diff --git a/tests/test_proto/utils/test_escape.py b/tests/test_proto/utils/test_escape.py new file mode 100644 index 00000000..1e618301 --- /dev/null +++ b/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