From 69d1a14b497b5a4d2d71eaea7b08e7eb0dc4718f Mon Sep 17 00:00:00 2001 From: Bilal Tonga Date: Mon, 19 Oct 2026 11:06:05 +0200 Subject: [PATCH 1/5] Refactor: Split Cursor into a shared CursorBase and the sync Cursor --- src/crate/client/cursor.py | 223 +++++++++++++++++++++++-------------- 1 file changed, 139 insertions(+), 84 deletions(-) diff --git a/src/crate/client/cursor.py b/src/crate/client/cursor.py index 0e331223..cf1d1974 100644 --- a/src/crate/client/cursor.py +++ b/src/crate/client/cursor.py @@ -113,10 +113,12 @@ def _convert_named_bulk_params( return converted_sql, bulk_args -class Cursor: +class CursorBase: """ - not thread-safe by intention - should not be shared between different threads + Cursor state and result handling shared by the sync and async cursors. + + Public methods that do IO, or that are coroutines on the async cursor, + live on the subclasses and call the private helpers here. """ lastrowid = None # currently not supported @@ -131,9 +133,11 @@ def __init__(self, connection, converter: Converter, **kwargs): self._time_zone = None self.time_zone = kwargs.get("time_zone") - def execute(self, sql, parameters=None, bulk_parameters=None): + def _prepare(self, sql, parameters=None, bulk_parameters=None): """ - Prepare and execute a database operation (query or command). + Check that the cursor can execute, and rewrite named parameters + to positional ones. Return the `sql`, `parameters` and + `bulk_parameters` to send. """ if self.connection._closed: raise ProgrammingError("Connection closed") @@ -150,24 +154,24 @@ def execute(self, sql, parameters=None, bulk_parameters=None): ) else: sql = _rewrite_pyformat_sql(sql) + return sql, parameters, bulk_parameters - self._result = self.connection.client.sql( - sql, parameters, bulk_parameters - ) + def _set_result(self, result): + """ + Store the response of a request and make its rows available. + """ + self._result = result if "rows" in self._result: if self._converter is None: self.rows = iter(self._result["rows"]) else: self.rows = iter(self._convert_rows()) - def executemany(self, sql, seq_of_parameters): + def _prepare_many(self, sql, seq_of_parameters): """ - Prepare a database operation (query or command) and then execute it - against all parameter sequences or mappings found in the sequence - ``seq_of_parameters``. + Rewrite a sequence of named parameter dicts to positional bulk + parameters. Return the `sql` and `bulk_parameters` to send. """ - row_counts = [] - durations = [] bulk_parameters = seq_of_parameters if ( bulk_parameters @@ -177,87 +181,78 @@ def executemany(self, sql, seq_of_parameters): sql, bulk_parameters = _convert_named_bulk_params( sql, bulk_parameters ) - self.execute(sql, bulk_parameters=bulk_parameters) + return sql, bulk_parameters + def _bulk_result(self): + """ + Replace the result of a bulk request with the aggregated row count + and duration, and return the per-statement results. + """ + row_counts = [] + durations = [] for result in self._result.get("results", []): if result.get("rowcount") > -1: row_counts.append(result.get("rowcount")) if self.duration > -1: durations.append(self.duration) - self._result = { - "rowcount": sum(row_counts) if row_counts else -1, - "duration": sum(durations) if durations else -1, - "rows": [], - "cols": self._result.get("cols", []), - "col_types": self._result.get("col_types", []), - "results": self._result.get("results"), - } - if self._converter is None: - self.rows = iter(self._result["rows"]) - else: - self.rows = iter(self._convert_rows()) + self._set_result( + { + "rowcount": sum(row_counts) if row_counts else -1, + "duration": sum(durations) if durations else -1, + "rows": [], + "cols": self._result.get("cols", []), + "col_types": self._result.get("col_types", []), + "results": self._result.get("results"), + } + ) return self._result["results"] - def fetchone(self): + def _next_row(self): """ - Fetch the next row of a query result set, returning a single sequence, - or None when no more data is available. - Alias for ``next()``. + Return the next row of a query result set, respecting if cursor was + closed. """ + if self.rows is None: + raise ProgrammingError( + "No result available. " + + "execute() or executemany() must be called first." + ) + if not self._closed: + return next(self.rows) + else: + raise ProgrammingError("Cursor closed") + + def _fetchone(self): try: - return self.next() + return self._next_row() except StopIteration: return None - def __iter__(self): - """ - support iterator interface: - http://legacy.python.org/dev/peps/pep-0249/#iter - - This iterator is shared. Advancing this iterator will advance other - iterators created from this cursor. - """ - warnings.warn("DB-API extension cursor.__iter__() used", stacklevel=2) - return self - - def fetchmany(self, count=None): - """ - Fetch the next set of rows of a query result, returning a sequence of - sequences (e.g. a list of tuples). An empty sequence is returned when - no more rows are available. - """ - if count is None: - count = self.arraysize - if count == 0: - return self.fetchall() + def _fetchmany(self, size=None): + if size is None: + size = self.arraysize + if size == 0: + return self._fetchall() result = [] - for _ in range(count): + for _ in range(size): try: - result.append(self.next()) + result.append(self._next_row()) except StopIteration: pass return result - def fetchall(self): - """ - Fetch all (remaining) rows of a query result, returning them as a - sequence of sequences (e.g. a list of tuples). Note that the cursor's - arraysize attribute can affect the performance of this operation. - """ + def _fetchall(self): result = [] iterate = True while iterate: try: - result.append(self.next()) + result.append(self._next_row()) except StopIteration: iterate = False return result - def close(self): - """ - Close the cursor now - """ + def _close(self): self._closed = True self._result = {} @@ -284,24 +279,6 @@ def rowcount(self): return -1 return self._result.get("rowcount", -1) - def next(self): - """ - Return the next row of a query result set, respecting if cursor was - closed. - """ - if self.rows is None: - raise ProgrammingError( - "No result available. " - + "execute() or executemany() must be called first." - ) - if not self._closed: - return next(self.rows) - else: - raise ProgrammingError("Cursor closed") - - def __next__(self): - return self.next() - @property def description(self): """ @@ -429,3 +406,81 @@ def _timezone_from_utc_offset(tz) -> timezone: raise ValueError( f"Time zone '{tz}' is given in invalid UTC offset format: {ex}" ) from ex + + +class Cursor(CursorBase): + """ + not thread-safe by intention + should not be shared between different threads + """ + + def execute(self, sql, parameters=None, bulk_parameters=None): + """ + Prepare and execute a database operation (query or command). + """ + self._set_result( + self.connection.client.sql( + *self._prepare(sql, parameters, bulk_parameters) + ) + ) + + def executemany(self, sql, seq_of_parameters): + """ + Prepare a database operation (query or command) and then execute it + against all parameter sequences or mappings found in the sequence + ``seq_of_parameters``. + """ + sql, bulk_parameters = self._prepare_many(sql, seq_of_parameters) + self.execute(sql, bulk_parameters=bulk_parameters) + return self._bulk_result() + + def fetchone(self): + """ + Fetch the next row of a query result set, returning a single sequence, + or None when no more data is available. + Alias for ``next()``. + """ + return self._fetchone() + + def __iter__(self): + """ + support iterator interface: + http://legacy.python.org/dev/peps/pep-0249/#iter + + This iterator is shared. Advancing this iterator will advance other + iterators created from this cursor. + """ + warnings.warn("DB-API extension cursor.__iter__() used", stacklevel=2) + return self + + def fetchmany(self, count=None): + """ + Fetch the next set of rows of a query result, returning a sequence of + sequences (e.g. a list of tuples). An empty sequence is returned when + no more rows are available. + """ + return self._fetchmany(count) + + def fetchall(self): + """ + Fetch all (remaining) rows of a query result, returning them as a + sequence of sequences (e.g. a list of tuples). Note that the cursor's + arraysize attribute can affect the performance of this operation. + """ + return self._fetchall() + + def close(self): + """ + Close the cursor now + """ + self._close() + + def next(self): + """ + Return the next row of a query result set, respecting if cursor was + closed. + """ + return self._next_row() + + def __next__(self): + return self.next() From 6e33ea5502043a4bc9e16045053ac86f480c6a1d Mon Sep 17 00:00:00 2001 From: Bilal Tonga Date: Mon, 19 Oct 2026 11:10:51 +0200 Subject: [PATCH 2/5] Refactor: Move CursorBase and the parameter helpers into _cursor_base.py --- src/crate/client/_cursor_base.py | 407 +++++++++++++++++++++++++++++++ src/crate/client/cursor.py | 393 +---------------------------- 2 files changed, 414 insertions(+), 386 deletions(-) create mode 100644 src/crate/client/_cursor_base.py diff --git a/src/crate/client/_cursor_base.py b/src/crate/client/_cursor_base.py new file mode 100644 index 00000000..d4f0635e --- /dev/null +++ b/src/crate/client/_cursor_base.py @@ -0,0 +1,407 @@ +# -*- coding: utf-8; -*- +# +# Licensed to CRATE Technology GmbH ("Crate") under one or more contributor +# license agreements. See the NOTICE file distributed with this work for +# additional information regarding copyright ownership. Crate licenses +# this file to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. You may +# obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations +# under the License. +# +# However, if you have executed another commercial license agreement +# with Crate these terms will supersede the license and you may use the +# software solely pursuant to the terms of the relevant commercial agreement. +import re +import typing as t +from datetime import datetime, timedelta, timezone +from itertools import count + +from .converter import Converter, DataType +from .exceptions import ProgrammingError + +_NAMED_PARAM_RE = re.compile(r"%\(([^)]+)\)s") + + +def _rewrite_pyformat_sql(sql: str) -> str: + """Replace %(name)s placeholders with $N positional markers (1-indexed).""" + counter = count(1) + return _NAMED_PARAM_RE.sub(lambda _: f"${next(counter)}", sql) + + +def _convert_named_to_positional( + sql: str, params: t.Dict[str, t.Any] +) -> t.Tuple[str, t.List[t.Any]]: + """Convert pyformat-style named parameters to positional parameters. + + Converts ``%(name)s`` placeholders to ``$N`` (1-indexed) and returns an + ordered list of corresponding values extracted from ``params``. + + The same name may appear multiple times; each occurrence appends the + value to the positional list independently. + + Raises ``ProgrammingError`` if a placeholder name is absent from ``params``. + Extra keys in ``params`` are silently ignored. + + Example:: + + sql = "SELECT * FROM t WHERE a = %(a)s AND b = %(b)s" + params = {"a": 1, "b": 2} + # returns: ("SELECT * FROM t WHERE a = $1 AND b = $2", [1, 2]) + """ + positions = {} + idx = 1 + new_params = [] + for k, v in params.items(): + positions[k] = idx + new_params.append(v) + idx += 1 + + def _replace(match: "re.Match[str]") -> str: + name = match.group(1) + if name not in params: + raise ProgrammingError( + f"Named parameter '{name}' not found in the parameters dict" + ) + position = positions[name] + return f"${position}" + + converted_sql = _NAMED_PARAM_RE.sub(_replace, sql) + return converted_sql, new_params + + +def _convert_named_bulk_params( + sql: str, seq_of_dicts: t.Sequence[t.Dict[str, t.Any]] +) -> t.Tuple[str, t.List[t.List[t.Any]]]: + """Convert pyformat SQL and a sequence of dicts to positional bulk args. + + Uses the first row to determine the SQL template and position map, then + builds a positional argument list for every row. + + Raises ``ProgrammingError`` if a placeholder name is absent from any row. + Extra keys in each row are silently ignored (consistent with + ``_convert_named_to_positional``). + """ + first = seq_of_dicts[0] + converted_sql, _ = _convert_named_to_positional(sql, first) + positions = {k: i + 1 for i, k in enumerate(first)} + n = len(positions) + + bulk_args: t.List[t.List[t.Any]] = [] + for row in seq_of_dicts: + if not isinstance(row, dict): + raise ProgrammingError( + "All bulk parameter rows must be dicts when SQL uses " + "pyformat (%(name)s) placeholders; got a non-dict row" + ) + positional: t.List[t.Any] = [None] * n + for name, pos in positions.items(): + if name not in row: + raise ProgrammingError( + f"Named parameter '{name}' not found in the parameters dict" + ) + positional[pos - 1] = row[name] + bulk_args.append(positional) + + return converted_sql, bulk_args + + +class CursorBase: + """ + Cursor state and result handling shared by the sync and async cursors. + + Public methods that do IO, or that are coroutines on the async cursor, + live on the subclasses and call the private helpers here. + """ + + lastrowid = None # currently not supported + + def __init__(self, connection, converter: Converter, **kwargs): + self.arraysize = 1 + self.connection = connection + self._converter = converter + self._closed = False + self._result: t.Dict[str, t.Any] = {} + self.rows = None + self._time_zone = None + self.time_zone = kwargs.get("time_zone") + + def _prepare(self, sql, parameters=None, bulk_parameters=None): + """ + Check that the cursor can execute, and rewrite named parameters + to positional ones. Return the `sql`, `parameters` and + `bulk_parameters` to send. + """ + if self.connection._closed: + raise ProgrammingError("Connection closed") + + if self._closed: + raise ProgrammingError("Cursor closed") + + if isinstance(parameters, dict): + sql, parameters = _convert_named_to_positional(sql, parameters) + elif bulk_parameters is not None and _NAMED_PARAM_RE.search(sql): + if bulk_parameters and isinstance(bulk_parameters[0], dict): + sql, bulk_parameters = _convert_named_bulk_params( + sql, bulk_parameters + ) + else: + sql = _rewrite_pyformat_sql(sql) + return sql, parameters, bulk_parameters + + def _set_result(self, result): + """ + Store the response of a request and make its rows available. + """ + self._result = result + if "rows" in self._result: + if self._converter is None: + self.rows = iter(self._result["rows"]) + else: + self.rows = iter(self._convert_rows()) + + def _prepare_many(self, sql, seq_of_parameters): + """ + Rewrite a sequence of named parameter dicts to positional bulk + parameters. Return the `sql` and `bulk_parameters` to send. + """ + bulk_parameters = seq_of_parameters + if ( + bulk_parameters + and isinstance(bulk_parameters[0], dict) + and _NAMED_PARAM_RE.search(sql) + ): + sql, bulk_parameters = _convert_named_bulk_params( + sql, bulk_parameters + ) + return sql, bulk_parameters + + def _bulk_result(self): + """ + Replace the result of a bulk request with the aggregated row count + and duration, and return the per-statement results. + """ + row_counts = [] + durations = [] + for result in self._result.get("results", []): + if result.get("rowcount") > -1: + row_counts.append(result.get("rowcount")) + if self.duration > -1: + durations.append(self.duration) + + self._set_result( + { + "rowcount": sum(row_counts) if row_counts else -1, + "duration": sum(durations) if durations else -1, + "rows": [], + "cols": self._result.get("cols", []), + "col_types": self._result.get("col_types", []), + "results": self._result.get("results"), + } + ) + return self._result["results"] + + def _next_row(self): + """ + Return the next row of a query result set, respecting if cursor was + closed. + """ + if self.rows is None: + raise ProgrammingError( + "No result available. " + + "execute() or executemany() must be called first." + ) + if not self._closed: + return next(self.rows) + else: + raise ProgrammingError("Cursor closed") + + def _fetchone(self): + try: + return self._next_row() + except StopIteration: + return None + + def _fetchmany(self, size=None): + if size is None: + size = self.arraysize + if size == 0: + return self._fetchall() + result = [] + for _ in range(size): + try: + result.append(self._next_row()) + except StopIteration: + pass + return result + + def _fetchall(self): + result = [] + iterate = True + while iterate: + try: + result.append(self._next_row()) + except StopIteration: + iterate = False + return result + + def _close(self): + self._closed = True + self._result = {} + + def setinputsizes(self, sizes): + """ + Not supported method. + """ + pass + + def setoutputsize(self, size, column=None): + """ + Not supported method. + """ + pass + + @property + def rowcount(self): + """ + This read-only attribute specifies the number of rows that the last + .execute*() produced (for DQL statements like ``SELECT``) or affected + (for DML statements like ``UPDATE`` or ``INSERT``). + """ + if self._closed or not self._result or "rows" not in self._result: + return -1 + return self._result.get("rowcount", -1) + + @property + def description(self): + """ + This read-only attribute is a sequence of 7-item sequences. + """ + if self._closed: + return None + + description = [] + for col in self._result["cols"]: + description.append((col, None, None, None, None, None, None)) + return tuple(description) + + @property + def duration(self): + """ + This read-only attribute specifies the server-side duration of a query + in milliseconds. + """ + if self._closed or not self._result or "duration" not in self._result: + return -1 + return self._result.get("duration", 0) + + def _convert_rows(self): + """ + Iterate rows, apply type converters, and generate converted rows. + """ + if not ("col_types" in self._result and self._result["col_types"]): + raise ValueError( + "Unable to apply type conversion " + "without `col_types` information" + ) + + # Resolve `col_types` definition to converter functions. Running + # the lookup redundantly on each row loop iteration would be a + # huge performance hog. + types = self._result["col_types"] + converters = [self._converter.get(type_) for type_ in types] + + # Process result rows with conversion. + for row in self._result["rows"]: + yield [ + convert(value) + for convert, value in zip(converters, row, strict=False) + ] + + @property + def time_zone(self): + """ + Get the current time zone. + """ + return self._time_zone + + @time_zone.setter + def time_zone(self, tz): + """ + Set the time zone. + + Different data types are supported. Available options are: + + - ``datetime.timezone.utc`` + - ``datetime.timezone(datetime.timedelta(hours=7), name="MST")`` + - ``pytz.timezone("Australia/Sydney")`` + - ``zoneinfo.ZoneInfo("Australia/Sydney")`` + - ``+0530`` (UTC offset in string format) + + The driver always returns timezone-"aware" `datetime` objects, + with their `tzinfo` attribute set. + + When `time_zone` is `None`, the returned `datetime` objects are + using Coordinated Universal Time (UTC), because CrateDB is storing + timestamp values in this format. + + When `time_zone` is given, the timestamp values will be transparently + converted from UTC to use the given time zone. + """ + + # Do nothing when time zone is reset. + if tz is None: + self._time_zone = None + return + + # Requesting datetime-aware `datetime` objects + # needs the data type converter. + # Implicitly create one, when needed. + if self._converter is None: + self._converter = Converter() + + # When the time zone is given as a string, + # assume UTC offset format, e.g. `+0530`. + if isinstance(tz, str): + tz = self._timezone_from_utc_offset(tz) + + self._time_zone = tz + + def _to_datetime_with_tz( + value: t.Optional[float], + ) -> t.Optional[datetime]: + """ + Convert CrateDB's `TIMESTAMP` value to a native Python `datetime` + object, with timezone-awareness. + """ + if value is None: + return None + return datetime.fromtimestamp(value / 1e3, tz=self._time_zone) + + # Register converter function for `TIMESTAMP` type. + self._converter.set(DataType.TIMESTAMP_WITH_TZ, _to_datetime_with_tz) + self._converter.set(DataType.TIMESTAMP_WITHOUT_TZ, _to_datetime_with_tz) + + @staticmethod + def _timezone_from_utc_offset(tz) -> timezone: + """ + UTC offset in string format (e.g. `+0530`) to `datetime.timezone`. + """ + if len(tz) != 5: + raise ValueError( + f"Time zone '{tz}' is given in invalid UTC offset format" + ) + try: + hours = int(tz[:3]) + minutes = int(tz[0] + tz[3:]) + return timezone(timedelta(hours=hours, minutes=minutes), name=tz) + except Exception as ex: + raise ValueError( + f"Time zone '{tz}' is given in invalid UTC offset format: {ex}" + ) from ex diff --git a/src/crate/client/cursor.py b/src/crate/client/cursor.py index cf1d1974..1433db0f 100644 --- a/src/crate/client/cursor.py +++ b/src/crate/client/cursor.py @@ -18,394 +18,15 @@ # However, if you have executed another commercial license agreement # with Crate these terms will supersede the license and you may use the # software solely pursuant to the terms of the relevant commercial agreement. -import re -import typing as t import warnings -from datetime import datetime, timedelta, timezone -from itertools import count -from .converter import Converter, DataType -from .exceptions import ProgrammingError - -_NAMED_PARAM_RE = re.compile(r"%\(([^)]+)\)s") - - -def _rewrite_pyformat_sql(sql: str) -> str: - """Replace %(name)s placeholders with $N positional markers (1-indexed).""" - counter = count(1) - return _NAMED_PARAM_RE.sub(lambda _: f"${next(counter)}", sql) - - -def _convert_named_to_positional( - sql: str, params: t.Dict[str, t.Any] -) -> t.Tuple[str, t.List[t.Any]]: - """Convert pyformat-style named parameters to positional parameters. - - Converts ``%(name)s`` placeholders to ``$N`` (1-indexed) and returns an - ordered list of corresponding values extracted from ``params``. - - The same name may appear multiple times; each occurrence appends the - value to the positional list independently. - - Raises ``ProgrammingError`` if a placeholder name is absent from ``params``. - Extra keys in ``params`` are silently ignored. - - Example:: - - sql = "SELECT * FROM t WHERE a = %(a)s AND b = %(b)s" - params = {"a": 1, "b": 2} - # returns: ("SELECT * FROM t WHERE a = $1 AND b = $2", [1, 2]) - """ - positions = {} - idx = 1 - new_params = [] - for k, v in params.items(): - positions[k] = idx - new_params.append(v) - idx += 1 - - def _replace(match: "re.Match[str]") -> str: - name = match.group(1) - if name not in params: - raise ProgrammingError( - f"Named parameter '{name}' not found in the parameters dict" - ) - position = positions[name] - return f"${position}" - - converted_sql = _NAMED_PARAM_RE.sub(_replace, sql) - return converted_sql, new_params - - -def _convert_named_bulk_params( - sql: str, seq_of_dicts: t.Sequence[t.Dict[str, t.Any]] -) -> t.Tuple[str, t.List[t.List[t.Any]]]: - """Convert pyformat SQL and a sequence of dicts to positional bulk args. - - Uses the first row to determine the SQL template and position map, then - builds a positional argument list for every row. - - Raises ``ProgrammingError`` if a placeholder name is absent from any row. - Extra keys in each row are silently ignored (consistent with - ``_convert_named_to_positional``). - """ - first = seq_of_dicts[0] - converted_sql, _ = _convert_named_to_positional(sql, first) - positions = {k: i + 1 for i, k in enumerate(first)} - n = len(positions) - - bulk_args: t.List[t.List[t.Any]] = [] - for row in seq_of_dicts: - if not isinstance(row, dict): - raise ProgrammingError( - "All bulk parameter rows must be dicts when SQL uses " - "pyformat (%(name)s) placeholders; got a non-dict row" - ) - positional: t.List[t.Any] = [None] * n - for name, pos in positions.items(): - if name not in row: - raise ProgrammingError( - f"Named parameter '{name}' not found in the parameters dict" - ) - positional[pos - 1] = row[name] - bulk_args.append(positional) - - return converted_sql, bulk_args - - -class CursorBase: - """ - Cursor state and result handling shared by the sync and async cursors. - - Public methods that do IO, or that are coroutines on the async cursor, - live on the subclasses and call the private helpers here. - """ - - lastrowid = None # currently not supported - - def __init__(self, connection, converter: Converter, **kwargs): - self.arraysize = 1 - self.connection = connection - self._converter = converter - self._closed = False - self._result: t.Dict[str, t.Any] = {} - self.rows = None - self._time_zone = None - self.time_zone = kwargs.get("time_zone") - - def _prepare(self, sql, parameters=None, bulk_parameters=None): - """ - Check that the cursor can execute, and rewrite named parameters - to positional ones. Return the `sql`, `parameters` and - `bulk_parameters` to send. - """ - if self.connection._closed: - raise ProgrammingError("Connection closed") - - if self._closed: - raise ProgrammingError("Cursor closed") - - if isinstance(parameters, dict): - sql, parameters = _convert_named_to_positional(sql, parameters) - elif bulk_parameters is not None and _NAMED_PARAM_RE.search(sql): - if bulk_parameters and isinstance(bulk_parameters[0], dict): - sql, bulk_parameters = _convert_named_bulk_params( - sql, bulk_parameters - ) - else: - sql = _rewrite_pyformat_sql(sql) - return sql, parameters, bulk_parameters - - def _set_result(self, result): - """ - Store the response of a request and make its rows available. - """ - self._result = result - if "rows" in self._result: - if self._converter is None: - self.rows = iter(self._result["rows"]) - else: - self.rows = iter(self._convert_rows()) - - def _prepare_many(self, sql, seq_of_parameters): - """ - Rewrite a sequence of named parameter dicts to positional bulk - parameters. Return the `sql` and `bulk_parameters` to send. - """ - bulk_parameters = seq_of_parameters - if ( - bulk_parameters - and isinstance(bulk_parameters[0], dict) - and _NAMED_PARAM_RE.search(sql) - ): - sql, bulk_parameters = _convert_named_bulk_params( - sql, bulk_parameters - ) - return sql, bulk_parameters - - def _bulk_result(self): - """ - Replace the result of a bulk request with the aggregated row count - and duration, and return the per-statement results. - """ - row_counts = [] - durations = [] - for result in self._result.get("results", []): - if result.get("rowcount") > -1: - row_counts.append(result.get("rowcount")) - if self.duration > -1: - durations.append(self.duration) - - self._set_result( - { - "rowcount": sum(row_counts) if row_counts else -1, - "duration": sum(durations) if durations else -1, - "rows": [], - "cols": self._result.get("cols", []), - "col_types": self._result.get("col_types", []), - "results": self._result.get("results"), - } - ) - return self._result["results"] - - def _next_row(self): - """ - Return the next row of a query result set, respecting if cursor was - closed. - """ - if self.rows is None: - raise ProgrammingError( - "No result available. " - + "execute() or executemany() must be called first." - ) - if not self._closed: - return next(self.rows) - else: - raise ProgrammingError("Cursor closed") - - def _fetchone(self): - try: - return self._next_row() - except StopIteration: - return None - - def _fetchmany(self, size=None): - if size is None: - size = self.arraysize - if size == 0: - return self._fetchall() - result = [] - for _ in range(size): - try: - result.append(self._next_row()) - except StopIteration: - pass - return result - - def _fetchall(self): - result = [] - iterate = True - while iterate: - try: - result.append(self._next_row()) - except StopIteration: - iterate = False - return result - - def _close(self): - self._closed = True - self._result = {} - - def setinputsizes(self, sizes): - """ - Not supported method. - """ - pass - - def setoutputsize(self, size, column=None): - """ - Not supported method. - """ - pass - - @property - def rowcount(self): - """ - This read-only attribute specifies the number of rows that the last - .execute*() produced (for DQL statements like ``SELECT``) or affected - (for DML statements like ``UPDATE`` or ``INSERT``). - """ - if self._closed or not self._result or "rows" not in self._result: - return -1 - return self._result.get("rowcount", -1) - - @property - def description(self): - """ - This read-only attribute is a sequence of 7-item sequences. - """ - if self._closed: - return None - - description = [] - for col in self._result["cols"]: - description.append((col, None, None, None, None, None, None)) - return tuple(description) - - @property - def duration(self): - """ - This read-only attribute specifies the server-side duration of a query - in milliseconds. - """ - if self._closed or not self._result or "duration" not in self._result: - return -1 - return self._result.get("duration", 0) - - def _convert_rows(self): - """ - Iterate rows, apply type converters, and generate converted rows. - """ - if not ("col_types" in self._result and self._result["col_types"]): - raise ValueError( - "Unable to apply type conversion " - "without `col_types` information" - ) - - # Resolve `col_types` definition to converter functions. Running - # the lookup redundantly on each row loop iteration would be a - # huge performance hog. - types = self._result["col_types"] - converters = [self._converter.get(type_) for type_ in types] - - # Process result rows with conversion. - for row in self._result["rows"]: - yield [ - convert(value) - for convert, value in zip(converters, row, strict=False) - ] - - @property - def time_zone(self): - """ - Get the current time zone. - """ - return self._time_zone - - @time_zone.setter - def time_zone(self, tz): - """ - Set the time zone. - - Different data types are supported. Available options are: - - - ``datetime.timezone.utc`` - - ``datetime.timezone(datetime.timedelta(hours=7), name="MST")`` - - ``pytz.timezone("Australia/Sydney")`` - - ``zoneinfo.ZoneInfo("Australia/Sydney")`` - - ``+0530`` (UTC offset in string format) - - The driver always returns timezone-"aware" `datetime` objects, - with their `tzinfo` attribute set. - - When `time_zone` is `None`, the returned `datetime` objects are - using Coordinated Universal Time (UTC), because CrateDB is storing - timestamp values in this format. - - When `time_zone` is given, the timestamp values will be transparently - converted from UTC to use the given time zone. - """ - - # Do nothing when time zone is reset. - if tz is None: - self._time_zone = None - return - - # Requesting datetime-aware `datetime` objects - # needs the data type converter. - # Implicitly create one, when needed. - if self._converter is None: - self._converter = Converter() - - # When the time zone is given as a string, - # assume UTC offset format, e.g. `+0530`. - if isinstance(tz, str): - tz = self._timezone_from_utc_offset(tz) - - self._time_zone = tz - - def _to_datetime_with_tz( - value: t.Optional[float], - ) -> t.Optional[datetime]: - """ - Convert CrateDB's `TIMESTAMP` value to a native Python `datetime` - object, with timezone-awareness. - """ - if value is None: - return None - return datetime.fromtimestamp(value / 1e3, tz=self._time_zone) - - # Register converter function for `TIMESTAMP` type. - self._converter.set(DataType.TIMESTAMP_WITH_TZ, _to_datetime_with_tz) - self._converter.set(DataType.TIMESTAMP_WITHOUT_TZ, _to_datetime_with_tz) - - @staticmethod - def _timezone_from_utc_offset(tz) -> timezone: - """ - UTC offset in string format (e.g. `+0530`) to `datetime.timezone`. - """ - if len(tz) != 5: - raise ValueError( - f"Time zone '{tz}' is given in invalid UTC offset format" - ) - try: - hours = int(tz[:3]) - minutes = int(tz[0] + tz[3:]) - return timezone(timedelta(hours=hours, minutes=minutes), name=tz) - except Exception as ex: - raise ValueError( - f"Time zone '{tz}' is given in invalid UTC offset format: {ex}" - ) from ex +from ._cursor_base import ( # noqa: F401 + _NAMED_PARAM_RE, + CursorBase, + _convert_named_bulk_params, + _convert_named_to_positional, + _rewrite_pyformat_sql, +) class Cursor(CursorBase): From e4e0838b3b866ed71fc8a8f37857b92576bc270b Mon Sep 17 00:00:00 2001 From: Bilal Tonga Date: Mon, 19 Oct 2026 11:15:32 +0200 Subject: [PATCH 3/5] Refactor: Split Connection into a shared ConnectionBase and the sync Connection --- src/crate/client/connection.py | 145 +++++++++++++++++++++++--------- tests/client/test_connection.py | 16 ++++ 2 files changed, 119 insertions(+), 42 deletions(-) diff --git a/src/crate/client/connection.py b/src/crate/client/connection.py index e499ce14..06f1780a 100644 --- a/src/crate/client/connection.py +++ b/src/crate/client/connection.py @@ -18,6 +18,7 @@ # However, if you have executed another commercial license agreement # with Crate these terms will supersede the license and you may use the # software solely pursuant to the terms of the relevant commercial agreement. +import typing as t from typing import Union from verlib2 import Version @@ -28,7 +29,99 @@ from .http import Client -class Connection: +class ConnectionBase: + """ + Connection state shared by the sync and async connections. + + Public methods that do IO, or that are coroutines on the async + connection, live on the subclasses. + """ + + client: t.Any + + @property + def _cursor_class(self) -> t.Type[t.Any]: + """The class `cursor()` creates.""" + raise NotImplementedError + + @property + def _blob_container_class(self) -> t.Type[t.Any]: + """The class `get_blob_container()` creates.""" + raise NotImplementedError + + def __init__(self, converter=None, time_zone=None): + self._converter = converter + self.time_zone = time_zone + self._closed = False + + def cursor(self, **kwargs): + """ + Return a new Cursor Object using the connection. + """ + converter = kwargs.pop("converter", self._converter) + time_zone = kwargs.pop("time_zone", self.time_zone) + self._check_open() + return self._cursor_class( + connection=self, + converter=converter, + time_zone=time_zone, + ) + + def get_blob_container(self, container_name): + """Retrieve a BlobContainer for `container_name` + + :param container_name: the name of the BLOB container. + :returns: a :class:ContainerObject + """ + return self._blob_container_class(container_name, self) + + def _check_open(self): + if self._closed: + raise ProgrammingError("Connection closed") + + @staticmethod + def _lowest_version( + versions: t.Iterable[t.Union[str, BaseException]], + ) -> Version: + """ + Return the lowest of the versions reported by the servers. + """ + lowest = None + connection_errors = [] + for version in versions: + if isinstance(version, ConnectionError): + connection_errors.append(version) + continue + if isinstance(version, ValueError): + continue + if isinstance(version, BaseException): + raise version + try: + parsed = Version(version) + except ValueError: + continue + if not lowest or parsed < lowest: + lowest = parsed + if connection_errors and not lowest: + raise ConnectionError("; ".join(str(e) for e in connection_errors)) + return lowest or Version("0.0.0") + + def __repr__(self): + return f"<{self.__class__.__qualname__} {self.client!r}>" + + +class Connection(ConnectionBase): + # Look up the module globals on each call, as before the base class + # existed: tests, e.g. in sqlalchemy-cratedb, patch + # `crate.client.connection.Cursor`. + @property + def _cursor_class(self) -> t.Type[Cursor]: + return Cursor + + @property + def _blob_container_class(self) -> t.Type[BlobContainer]: + return BlobContainer + def __init__( self, servers=None, @@ -141,8 +234,7 @@ def __init__( An integer compresses only when the payload exceeds that many bytes. """ # noqa: E501 - self._converter = converter - self.time_zone = time_zone + super().__init__(converter=converter, time_zone=time_zone) if client: self.client = client @@ -169,22 +261,12 @@ def __init__( compress=compress, ) self.lowest_server_version = self._lowest_server_version() - self._closed = False def cursor(self, **kwargs) -> Cursor: """ Return a new Cursor Object using the connection. """ - converter = kwargs.pop("converter", self._converter) - time_zone = kwargs.pop("time_zone", self.time_zone) - if not self._closed: - return Cursor( - connection=self, - converter=converter, - time_zone=time_zone, - ) - else: - raise ProgrammingError("Connection closed") + return super().cursor(**kwargs) def close(self): """ @@ -197,38 +279,17 @@ def commit(self): """ Transactions are not supported, so ``commit`` is not implemented. """ - if self._closed: - raise ProgrammingError("Connection closed") - - def get_blob_container(self, container_name): - """Retrieve a BlobContainer for `container_name` - - :param container_name: the name of the BLOB container. - :returns: a :class:ContainerObject - """ - return BlobContainer(container_name, self) + self._check_open() def _lowest_server_version(self): - lowest = None - servers = self.client.active_servers - connection_errors = [] - for server in servers: + versions: t.List[t.Union[str, BaseException]] = [] + for server in self.client.active_servers: try: _, _, version = self.client.server_infos(server) - version = Version(version) - except ConnectionError as ex: - connection_errors.append(ex) - continue - except ValueError: - continue - if not lowest or version < lowest: - lowest = version - if connection_errors and not lowest: - raise ConnectionError("; ".join(str(e) for e in connection_errors)) - return lowest or Version("0.0.0") - - def __repr__(self): - return f"<{self.__class__.__qualname__} {self.client!r}>" + versions.append(version) + except (ConnectionError, ValueError) as ex: + versions.append(ex) + return self._lowest_version(versions) def __enter__(self): return self diff --git a/tests/client/test_connection.py b/tests/client/test_connection.py index 464ec123..1c7de512 100644 --- a/tests/client/test_connection.py +++ b/tests/client/test_connection.py @@ -317,3 +317,19 @@ def server_infos(self, server): # 3 != 2 → guard never fires → silently returns Version("0.0.0"). with pytest.raises(crate.client.exceptions.ConnectionError): conn._lowest_server_version() + + +def test_patched_cursor_and_blob_container_classes_are_used(): + """ + `Connection` looks up `Cursor` and `BlobContainer` in its module on + each call, so patching them there takes effect. sqlalchemy-cratedb's + tests patch `crate.client.connection.Cursor`. + """ + client = _FakeClient(["http://localhost:4200"], lambda _: (0, 0, "6.0.0")) + conn = connect(client=client) + with ( + patch("crate.client.connection.Cursor") as cursor_class, + patch("crate.client.connection.BlobContainer") as container_class, + ): + assert conn.cursor() is cursor_class.return_value + assert conn.get_blob_container("blobs") is container_class.return_value From 0b2c47f8abd8fbe5f45f9e02d8bad7c2e6ca8fa0 Mon Sep 17 00:00:00 2001 From: Bilal Tonga Date: Mon, 19 Oct 2026 11:18:39 +0200 Subject: [PATCH 4/5] Refactor: Move ConnectionBase into _connection_base.py --- pyproject.toml | 2 +- src/crate/client/_connection_base.py | 106 +++++++++++++++++++++++++++ src/crate/client/connection.py | 86 +--------------------- 3 files changed, 109 insertions(+), 85 deletions(-) create mode 100644 src/crate/client/_connection_base.py diff --git a/pyproject.toml b/pyproject.toml index 13bacc48..03b9a034 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -174,7 +174,7 @@ lint.per-file-ignores."tests/*" = [ "S310", # False positive; it's a https url "S311", # Standard pseudo-random generators are not suitable for cryptographic purposes ] -lint.per-file-ignores."src/crate/client/{connection.py,http.py,_pool.py,_response.py}" = [ +lint.per-file-ignores."src/crate/client/{connection.py,http.py,_connection_base.py,_pool.py,_response.py}" = [ "A004", # Import `ConnectionError` is shadowing a Python builtin "A005", # Import `ConnectionError` is shadowing a Python builtin ] diff --git a/src/crate/client/_connection_base.py b/src/crate/client/_connection_base.py new file mode 100644 index 00000000..2c9131b2 --- /dev/null +++ b/src/crate/client/_connection_base.py @@ -0,0 +1,106 @@ +# -*- coding: utf-8; -*- +# +# Licensed to CRATE Technology GmbH ("Crate") under one or more contributor +# license agreements. See the NOTICE file distributed with this work for +# additional information regarding copyright ownership. Crate licenses +# this file to you under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. You may +# obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations +# under the License. +# +# However, if you have executed another commercial license agreement +# with Crate these terms will supersede the license and you may use the +# software solely pursuant to the terms of the relevant commercial agreement. +import typing as t + +from verlib2 import Version + +from .exceptions import ConnectionError, ProgrammingError + + +class ConnectionBase: + """ + Connection state shared by the sync and async connections. + + Public methods that do IO, or that are coroutines on the async + connection, live on the subclasses. + """ + + client: t.Any + + @property + def _cursor_class(self) -> t.Type[t.Any]: + """The class `cursor()` creates.""" + raise NotImplementedError + + @property + def _blob_container_class(self) -> t.Type[t.Any]: + """The class `get_blob_container()` creates.""" + raise NotImplementedError + + def __init__(self, converter=None, time_zone=None): + self._converter = converter + self.time_zone = time_zone + self._closed = False + + def cursor(self, **kwargs): + """ + Return a new Cursor Object using the connection. + """ + converter = kwargs.pop("converter", self._converter) + time_zone = kwargs.pop("time_zone", self.time_zone) + self._check_open() + return self._cursor_class( + connection=self, + converter=converter, + time_zone=time_zone, + ) + + def get_blob_container(self, container_name): + """Retrieve a BlobContainer for `container_name` + + :param container_name: the name of the BLOB container. + :returns: a :class:ContainerObject + """ + return self._blob_container_class(container_name, self) + + def _check_open(self): + if self._closed: + raise ProgrammingError("Connection closed") + + @staticmethod + def _lowest_version( + versions: t.Iterable[t.Union[str, BaseException]], + ) -> Version: + """ + Return the lowest of the versions reported by the servers. + """ + lowest = None + connection_errors = [] + for version in versions: + if isinstance(version, ConnectionError): + connection_errors.append(version) + continue + if isinstance(version, ValueError): + continue + if isinstance(version, BaseException): + raise version + try: + parsed = Version(version) + except ValueError: + continue + if not lowest or parsed < lowest: + lowest = parsed + if connection_errors and not lowest: + raise ConnectionError("; ".join(str(e) for e in connection_errors)) + return lowest or Version("0.0.0") + + def __repr__(self): + return f"<{self.__class__.__qualname__} {self.client!r}>" diff --git a/src/crate/client/connection.py b/src/crate/client/connection.py index 06f1780a..12930c17 100644 --- a/src/crate/client/connection.py +++ b/src/crate/client/connection.py @@ -21,95 +21,13 @@ import typing as t from typing import Union -from verlib2 import Version - +from ._connection_base import ConnectionBase from .blob import BlobContainer from .cursor import Cursor -from .exceptions import ConnectionError, ProgrammingError +from .exceptions import ConnectionError from .http import Client -class ConnectionBase: - """ - Connection state shared by the sync and async connections. - - Public methods that do IO, or that are coroutines on the async - connection, live on the subclasses. - """ - - client: t.Any - - @property - def _cursor_class(self) -> t.Type[t.Any]: - """The class `cursor()` creates.""" - raise NotImplementedError - - @property - def _blob_container_class(self) -> t.Type[t.Any]: - """The class `get_blob_container()` creates.""" - raise NotImplementedError - - def __init__(self, converter=None, time_zone=None): - self._converter = converter - self.time_zone = time_zone - self._closed = False - - def cursor(self, **kwargs): - """ - Return a new Cursor Object using the connection. - """ - converter = kwargs.pop("converter", self._converter) - time_zone = kwargs.pop("time_zone", self.time_zone) - self._check_open() - return self._cursor_class( - connection=self, - converter=converter, - time_zone=time_zone, - ) - - def get_blob_container(self, container_name): - """Retrieve a BlobContainer for `container_name` - - :param container_name: the name of the BLOB container. - :returns: a :class:ContainerObject - """ - return self._blob_container_class(container_name, self) - - def _check_open(self): - if self._closed: - raise ProgrammingError("Connection closed") - - @staticmethod - def _lowest_version( - versions: t.Iterable[t.Union[str, BaseException]], - ) -> Version: - """ - Return the lowest of the versions reported by the servers. - """ - lowest = None - connection_errors = [] - for version in versions: - if isinstance(version, ConnectionError): - connection_errors.append(version) - continue - if isinstance(version, ValueError): - continue - if isinstance(version, BaseException): - raise version - try: - parsed = Version(version) - except ValueError: - continue - if not lowest or parsed < lowest: - lowest = parsed - if connection_errors and not lowest: - raise ConnectionError("; ".join(str(e) for e in connection_errors)) - return lowest or Version("0.0.0") - - def __repr__(self): - return f"<{self.__class__.__qualname__} {self.client!r}>" - - class Connection(ConnectionBase): # Look up the module globals on each call, as before the base class # existed: tests, e.g. in sqlalchemy-cratedb, patch From 4b2b863724c045b033f53cbfee42073d7006b3ec Mon Sep 17 00:00:00 2001 From: Bilal Tonga Date: Mon, 19 Oct 2026 11:26:27 +0200 Subject: [PATCH 5/5] Tests: Cover ConnectionBase._lowest_version directly --- tests/client/test_connection.py | 45 +++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/tests/client/test_connection.py b/tests/client/test_connection.py index 1c7de512..cdfe49ef 100644 --- a/tests/client/test_connection.py +++ b/tests/client/test_connection.py @@ -319,6 +319,51 @@ def server_infos(self, server): conn._lowest_server_version() +def test_lowest_version_picks_lowest(): + versions = ["5.5.2", "1.0.3", "6.0.0"] + assert (1, 0, 3) == Connection._lowest_version(versions).version + + +def test_lowest_version_skips_invalid_versions(): + versions = ["not a version", ValueError("bad"), "5.5.2"] + assert (5, 5, 2) == Connection._lowest_version(versions).version + + +def test_lowest_version_defaults_without_versions(): + assert (0, 0, 0) == Connection._lowest_version([]).version + assert (0, 0, 0) == Connection._lowest_version(["invalid"]).version + + +def test_lowest_version_ignores_unreachable_servers(): + versions = [ + crate.client.exceptions.ConnectionError("node-A down"), + "5.5.2", + ] + assert (5, 5, 2) == Connection._lowest_version(versions).version + + +def test_lowest_version_raises_when_no_server_answered(): + versions = [ + crate.client.exceptions.ConnectionError("node-A down"), + crate.client.exceptions.ConnectionError("node-B down"), + "invalid", + ] + with pytest.raises(crate.client.exceptions.ConnectionError) as excinfo: + Connection._lowest_version(versions) + assert "node-A down; node-B down" == str(excinfo.value) + + +def test_lowest_version_raises_other_exceptions(): + """ + An exception that is neither a ConnectionError nor a ValueError, + for example a failed authentication, is raised as is. + """ + error = ProgrammingError("authentication failed") + with pytest.raises(ProgrammingError) as excinfo: + Connection._lowest_version(["5.5.2", error]) + assert excinfo.value is error + + def test_patched_cursor_and_blob_container_classes_are_used(): """ `Connection` looks up `Cursor` and `BlobContainer` in its module on