From 69d1a14b497b5a4d2d71eaea7b08e7eb0dc4718f Mon Sep 17 00:00:00 2001 From: Bilal Tonga Date: Mon, 19 Oct 2026 11:06:05 +0200 Subject: [PATCH 1/8] 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/8] 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/8] 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/8] 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/8] 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 From 3d6c483bc8cbe73f9c369269cf782d72747b2ff0 Mon Sep 17 00:00:00 2001 From: Bilal Tonga Date: Tue, 27 Oct 2026 16:59:22 +0100 Subject: [PATCH 6/8] Refactor ConnectionBase and Connection classes to support generic cursor types --- src/crate/client/_connection_base.py | 9 ++++++--- src/crate/client/connection.py | 8 +------- 2 files changed, 7 insertions(+), 10 deletions(-) diff --git a/src/crate/client/_connection_base.py b/src/crate/client/_connection_base.py index 2c9131b2..7660ae84 100644 --- a/src/crate/client/_connection_base.py +++ b/src/crate/client/_connection_base.py @@ -22,10 +22,13 @@ from verlib2 import Version +from ._cursor_base import CursorBase from .exceptions import ConnectionError, ProgrammingError +CursorT = t.TypeVar("CursorT", bound=CursorBase) -class ConnectionBase: + +class ConnectionBase(t.Generic[CursorT]): """ Connection state shared by the sync and async connections. @@ -36,7 +39,7 @@ class ConnectionBase: client: t.Any @property - def _cursor_class(self) -> t.Type[t.Any]: + def _cursor_class(self) -> t.Type[CursorT]: """The class `cursor()` creates.""" raise NotImplementedError @@ -50,7 +53,7 @@ def __init__(self, converter=None, time_zone=None): self.time_zone = time_zone self._closed = False - def cursor(self, **kwargs): + def cursor(self, **kwargs) -> CursorT: """ Return a new Cursor Object using the connection. """ diff --git a/src/crate/client/connection.py b/src/crate/client/connection.py index 12930c17..0bea6317 100644 --- a/src/crate/client/connection.py +++ b/src/crate/client/connection.py @@ -28,7 +28,7 @@ from .http import Client -class Connection(ConnectionBase): +class Connection(ConnectionBase[Cursor]): # 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`. @@ -180,12 +180,6 @@ def __init__( ) self.lowest_server_version = self._lowest_server_version() - def cursor(self, **kwargs) -> Cursor: - """ - Return a new Cursor Object using the connection. - """ - return super().cursor(**kwargs) - def close(self): """ Close the connection now From f1443232470d857e2e922c8ef93fadac48607d0d Mon Sep 17 00:00:00 2001 From: Bilal Tonga Date: Tue, 27 Oct 2026 17:21:53 +0100 Subject: [PATCH 7/8] Add asynchronous connection and cursor implementations --- src/crate/client/aio/__init__.py | 42 +++++++ src/crate/client/aio/connection.py | 177 +++++++++++++++++++++++++++++ src/crate/client/aio/cursor.py | 111 ++++++++++++++++++ 3 files changed, 330 insertions(+) create mode 100644 src/crate/client/aio/connection.py create mode 100644 src/crate/client/aio/cursor.py diff --git a/src/crate/client/aio/__init__.py b/src/crate/client/aio/__init__.py index 30da9ed9..ef59aaf6 100644 --- a/src/crate/client/aio/__init__.py +++ b/src/crate/client/aio/__init__.py @@ -33,3 +33,45 @@ "crate.client.aio requires aiohttp. " 'Install it with: pip install "aiohttp>=3.12,<4"' ) from ex + +from crate.client.aio._http import AsyncClient +from crate.client.aio.connection import AsyncConnection, connect +from crate.client.aio.cursor import AsyncCursor +from crate.client.exceptions import ( + BlobLocationNotFoundException, + ConnectionError, + DatabaseError, + DataError, + DigestNotFoundException, + Error, + IntegrityError, + InterfaceError, + InternalError, + NotSupportedError, + OperationalError, + ProgrammingError, + Warning, +) + +paramstyle = "pyformat" + +__all__ = [ + "AsyncClient", + "AsyncConnection", + "AsyncCursor", + "BlobLocationNotFoundException", + "ConnectionError", + "DataError", + "DatabaseError", + "DigestNotFoundException", + "Error", + "IntegrityError", + "InterfaceError", + "InternalError", + "NotSupportedError", + "OperationalError", + "ProgrammingError", + "Warning", + "connect", + "paramstyle", +] diff --git a/src/crate/client/aio/connection.py b/src/crate/client/aio/connection.py new file mode 100644 index 00000000..1ea77eb6 --- /dev/null +++ b/src/crate/client/aio/connection.py @@ -0,0 +1,177 @@ +# -*- 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 asyncio +import typing as t +import warnings + +from verlib2 import Version + +from crate.client._connection_base import ConnectionBase +from crate.client.aio._http import AsyncClient +from crate.client.aio.cursor import AsyncCursor + + +class AsyncConnection(ConnectionBase[AsyncCursor]): + """ + Asynchronous connection. Experimental. + """ + + client: AsyncClient + lowest_server_version: t.Optional[Version] + + @property + def _cursor_class(self) -> t.Type[AsyncCursor]: + return AsyncCursor + + def __init__( + self, + client: AsyncClient, + converter: t.Any = None, + time_zone: t.Any = None, + ) -> None: + super().__init__(converter=converter, time_zone=time_zone) + self.client = client + self.lowest_server_version = None + + @property + def closed(self) -> bool: + return self._closed + + async def commit(self) -> None: + """ + Transactions are not supported, so ``commit`` is not implemented. + """ + self._check_open() + + async def rollback(self) -> None: + """ + Transactions are not supported, so ``rollback`` does nothing. + """ + + async def close(self) -> None: + """ + Wait for in-flight requests, then close all connections. + """ + if self._closed: + return + self._closed = True + await self.client.close() + + def terminate(self) -> None: + """ + Close all connections at once, without awaiting anything. + """ + self._closed = True + self.client.terminate() + + async def _lowest_server_version(self) -> Version: + results = await asyncio.gather( + *( + self.client.server_infos(server) + for server in self.client.active_servers + ), + return_exceptions=True, + ) + return self._lowest_version( + result if isinstance(result, BaseException) else result[2] + for result in results + ) + + async def __aenter__(self) -> "AsyncConnection": + return self + + async def __aexit__(self, *excs: object) -> None: + await self.close() + + def __del__(self) -> None: + if getattr(self, "_closed", True): + return + warnings.warn( + f"unclosed {self!r}; close it with `await connection.close()`", + ResourceWarning, + stacklevel=2, + source=self, + ) + # Also keeps the client from warning a second time. + self.terminate() + + +async def connect( + servers: t.Union[str, t.Iterable[str], None] = None, + timeout: t.Union[int, float, str, None] = None, + backoff_factor: float = 0, + client: t.Optional[AsyncClient] = None, + verify_ssl_cert: bool = True, + ca_cert: t.Optional[str] = None, + error_trace: bool = False, + cert_file: t.Optional[str] = None, + key_file: t.Optional[str] = None, + ssl_relax_minimum_version: bool = False, + username: t.Optional[str] = None, + password: t.Optional[str] = None, + schema: t.Optional[str] = None, + pool_size: t.Optional[int] = None, + socket_keepalive: bool = True, + socket_tcp_keepidle: t.Optional[int] = None, + socket_tcp_keepintvl: t.Optional[int] = None, + socket_tcp_keepcnt: t.Optional[int] = None, + converter: t.Any = None, + time_zone: t.Any = None, + jwt_token: t.Optional[str] = None, + compress: t.Union[int, bool] = 8192, +) -> AsyncConnection: + """ + Connect to CrateDB and ask every server for its version. + """ + if client is None: + client = AsyncClient( + servers, + timeout=timeout, + backoff_factor=backoff_factor, + verify_ssl_cert=verify_ssl_cert, + ca_cert=ca_cert, + error_trace=error_trace, + cert_file=cert_file, + key_file=key_file, + ssl_relax_minimum_version=ssl_relax_minimum_version, + username=username, + password=password, + schema=schema, + pool_size=pool_size, + socket_keepalive=socket_keepalive, + socket_tcp_keepidle=socket_tcp_keepidle, + socket_tcp_keepintvl=socket_tcp_keepintvl, + socket_tcp_keepcnt=socket_tcp_keepcnt, + jwt_token=jwt_token, + compress=compress, + ) + connection = AsyncConnection( + client, converter=converter, time_zone=time_zone + ) + try: + connection.lowest_server_version = ( + await connection._lowest_server_version() + ) + except BaseException: + connection.terminate() + raise + return connection diff --git a/src/crate/client/aio/cursor.py b/src/crate/client/aio/cursor.py new file mode 100644 index 00000000..dc831d99 --- /dev/null +++ b/src/crate/client/aio/cursor.py @@ -0,0 +1,111 @@ +# -*- 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 crate.client._cursor_base import CursorBase + +if t.TYPE_CHECKING: + from crate.client.aio.connection import AsyncConnection + + +class AsyncCursor(CursorBase): + """ + Asynchronous cursor. Experimental. + """ + + connection: "AsyncConnection" + + async def execute( + self, + sql: str, + parameters: t.Any = None, + bulk_parameters: t.Any = None, + ) -> None: + """ + Prepare and execute a database operation (query or command). + """ + self._set_result( + await self.connection.client.sql( + *self._prepare(sql, parameters, bulk_parameters) + ) + ) + + async def executemany( + self, sql: str, seq_of_parameters: t.Sequence[t.Any] + ) -> t.List[t.Dict[str, t.Any]]: + """ + 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) + await self.execute(sql, bulk_parameters=bulk_parameters) + return self._bulk_result() + + async def fetchone(self) -> t.Optional[t.List[t.Any]]: + """ + Fetch the next row of a query result set, returning a single sequence, + or None when no more data is available. + """ + return self._fetchone() + + async def fetchmany( + self, size: t.Optional[int] = None + ) -> t.List[t.List[t.Any]]: + """ + 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(size) + + async def fetchall(self) -> t.List[t.List[t.Any]]: + """ + 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() + + async def close(self) -> None: + """ + Close the cursor now + """ + self._close() + + def __aiter__(self) -> "AsyncCursor": + """ + Iterate the remaining rows with ``async for``. + """ + return self + + async def __anext__(self) -> t.List[t.Any]: + try: + return self._next_row() + except StopIteration: + raise StopAsyncIteration from None + + async def __aenter__(self) -> "AsyncCursor": + return self + + async def __aexit__(self, *excs: object) -> None: + await self.close() From faabc5f8198fc801c8ace3c53be2f5760257bce3 Mon Sep 17 00:00:00 2001 From: Bilal Tonga Date: Tue, 27 Oct 2026 17:43:02 +0100 Subject: [PATCH 8/8] Add tests for cursor and connection --- tests/client/test_aio_connection.py | 369 ++++++++++++++++++++ tests/client/test_aio_cursor.py | 509 ++++++++++++++++++++++++++++ 2 files changed, 878 insertions(+) create mode 100644 tests/client/test_aio_connection.py create mode 100644 tests/client/test_aio_cursor.py diff --git a/tests/client/test_aio_connection.py b/tests/client/test_aio_connection.py new file mode 100644 index 00000000..206cfa1c --- /dev/null +++ b/tests/client/test_aio_connection.py @@ -0,0 +1,369 @@ +# -*- 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 asyncio +import gc +import inspect +import json +import warnings + +import pytest + +pytest.importorskip("aiohttp") + +from aiohttp import web # noqa: E402 +from verlib2 import Version # noqa: E402 + +import crate.client # noqa: E402 +from crate.client import aio, exceptions # noqa: E402 +from crate.client.aio import ( # noqa: E402 + AsyncClient, + AsyncConnection, + AsyncCursor, +) +from crate.client.converter import DefaultTypeConverter # noqa: E402 + +from .aio_support import FakeCrate, unused_port # noqa: E402 + +pytestmark = [ + pytest.mark.timeout(30, method="thread"), + pytest.mark.filterwarnings("error::ResourceWarning"), + pytest.mark.filterwarnings( + "error::pytest.PytestUnraisableExceptionWarning" + ), +] + + +def crate_handler(version="6.1.2", result=None): + + async def handler(request): + if request.path == "/": + return web.json_response( + {"name": "crate-1", "version": {"number": version}} + ) + return web.json_response( + result or {"cols": ["name"], "rows": [["a"]], "rowcount": 1} + ) + + return handler + + +def dead_server(): + return f"http://127.0.0.1:{unused_port()}" + + +def sync_failure(*args, **kwargs): + try: + crate.client.connect(*args, **kwargs).close() + except exceptions.Error as ex: + return type(ex), str(ex) + raise AssertionError("connect() did not fail") + + +async def async_failure(*args, **kwargs): + try: + await (await aio.connect(*args, **kwargs)).close() + except exceptions.Error as ex: + return type(ex), str(ex) + raise AssertionError("connect() did not fail") + + +def test_exports(): + assert aio.paramstyle == crate.client.paramstyle == "pyformat" + for name in [ + "Error", + "Warning", + "InterfaceError", + "DatabaseError", + "InternalError", + "OperationalError", + "ProgrammingError", + "IntegrityError", + "DataError", + "NotSupportedError", + "ConnectionError", + "BlobLocationNotFoundException", + "DigestNotFoundException", + ]: + assert getattr(aio, name) is getattr(exceptions, name), name + assert aio.connect is aio.connection.connect + assert aio.AsyncConnection is AsyncConnection + assert aio.AsyncCursor is AsyncCursor + assert aio.AsyncClient is AsyncClient + + +def test_connect_signature_matches_sync(): + sync = inspect.signature(crate.client.connect).parameters + actual = inspect.signature(aio.connect).parameters + assert list(actual) == list(sync) + for name, parameter in actual.items(): + assert parameter.default == sync[name].default, name + assert inspect.iscoroutinefunction(aio.connect) + + +@pytest.mark.asyncio +async def test_connect_sets_lowest_server_version(fake_crate): + old = await fake_crate(crate_handler(version="5.10.3")) + new = await fake_crate(crate_handler(version="6.1.2")) + async with await aio.connect([new.url, old.url]) as conn: + assert conn.lowest_server_version == Version("5.10.3") + assert not conn.closed + assert conn.closed + + +@pytest.mark.asyncio +async def test_connect_probes_servers_concurrently(fake_crate): + both_probed = asyncio.Event() + probes = [] + + async def handler(request): + probes.append(request.url.port) + if len(probes) == 2: + both_probed.set() + await both_probed.wait() + return web.json_response({"version": {"number": "6.1.2"}}) + + first = await fake_crate(handler) + second = await fake_crate(handler) + conn = await asyncio.wait_for(aio.connect([first.url, second.url]), 5) + await conn.close() + assert sorted(probes) == sorted([first.port, second.port]) + + +@pytest.mark.asyncio +async def test_connect_with_a_dead_server(fake_crate): + live = await fake_crate(crate_handler(version="5.10.3")) + async with await aio.connect([dead_server(), live.url]) as conn: + assert conn.lowest_server_version == Version("5.10.3") + + +@pytest.mark.asyncio +async def test_connect_with_all_servers_dead(): + servers = [dead_server(), dead_server()] + expected = await asyncio.to_thread(sync_failure, servers) + client = AsyncClient(servers) + actual = await async_failure(client=client) + assert actual[0] is expected[0] is exceptions.ConnectionError + assert client.closed + + +@pytest.mark.asyncio +async def test_connect_with_wrong_password(fake_crate): + + async def unauthorized(request): + return web.json_response( + { + "error": { + "message": "AuthenticationException[password " + "authentication failed for user me]", + "code": 40100, + } + }, + status=401, + reason="Unauthorized", + ) + + crate = await fake_crate(unauthorized) + kwargs = {"username": "me", "password": "wrong"} + expected = await asyncio.to_thread(sync_failure, crate.url, **kwargs) + client = AsyncClient(crate.url, **kwargs) + actual = await async_failure(client=client) + assert actual == expected + assert actual[0] is exceptions.ProgrammingError + assert client.closed + + +@pytest.mark.asyncio +async def test_cancelled_connect_closes_the_client(fake_crate): + probing = asyncio.Event() + + async def slow(request): + probing.set() + await asyncio.sleep(10) + return web.json_response({}) + + crate = await fake_crate(slow) + client = AsyncClient(servers=crate.url) + task = asyncio.ensure_future(aio.connect(client=client)) + await probing.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert client.closed + assert not any(server.is_open for server in client.server_pool.values()) + assert client.active_servers == [crate.url] + + +@pytest.mark.asyncio +async def test_cursor(fake_crate): + crate = await fake_crate( + crate_handler( + result={ + "cols": ["name"], + "col_types": [4], + "rows": [["a"], ["b"], ["c"]], + "rowcount": 3, + } + ) + ) + async with await aio.connect(crate.url) as conn: + cursor = conn.cursor() + assert isinstance(cursor, AsyncCursor) + assert cursor.connection is conn + await cursor.execute("SELECT name FROM t WHERE id = %(id)s", {"id": 1}) + assert await cursor.fetchone() == ["a"] + assert [row async for row in cursor] == [["b"], ["c"]] + await cursor.executemany("SELECT name FROM t WHERE id = ?", [[1], [2]]) + _, body = crate.requests[-2] + assert json.loads(body) == { + "stmt": "SELECT name FROM t WHERE id = $1", + "args": [1], + } + _, body = crate.requests[-1] + assert json.loads(body) == { + "stmt": "SELECT name FROM t WHERE id = ?", + "bulk_args": [[1], [2]], + } + + +@pytest.mark.asyncio +async def test_cursor_inherits_converter_and_time_zone(fake_crate): + """ + As in sync, an omitted `converter` or `time_zone` is taken from the + connection, and an explicit `None` means none. + """ + crate = await fake_crate(crate_handler()) + converter = DefaultTypeConverter() + async with await aio.connect( + crate.url, converter=converter, time_zone="+0530" + ) as conn: + inherited = conn.cursor() + assert inherited._converter is converter + assert str(inherited.time_zone) == "+0530" + explicit = conn.cursor(converter=None, time_zone=None) + assert explicit._converter is None + assert explicit.time_zone is None + + +@pytest.mark.asyncio +async def test_lifecycle(fake_crate): + crate = await fake_crate(crate_handler()) + conn = await aio.connect(crate.url) + cursor = conn.cursor() + assert inspect.iscoroutinefunction(conn.commit) + await conn.commit() + await conn.rollback() + await conn.close() + assert conn.closed + assert conn.client.closed + await conn.close() + with pytest.raises(exceptions.ProgrammingError, match="Connection closed"): + await conn.commit() + with pytest.raises(exceptions.ProgrammingError, match="Connection closed"): + conn.cursor() + with pytest.raises(exceptions.ProgrammingError, match="Connection closed"): + await cursor.execute("SELECT 1") + await conn.rollback() + + +@pytest.mark.asyncio +async def test_close_waits_for_in_flight_requests(fake_crate): + executing = asyncio.Event() + release = asyncio.Event() + + async def handler(request): + if request.path == "/": + return web.json_response({"version": {"number": "6.1.2"}}) + executing.set() + await release.wait() + return web.json_response({"cols": [], "rows": [[]], "rowcount": 1}) + + crate = await fake_crate(handler) + conn = await aio.connect(crate.url) + cursor = conn.cursor() + execute = asyncio.ensure_future(cursor.execute("INSERT INTO t VALUES (1)")) + await executing.wait() + close = asyncio.ensure_future(conn.close()) + await asyncio.sleep(0.1) + assert not close.done() + release.set() + await asyncio.wait_for(asyncio.gather(execute, close), 5) + assert cursor.rowcount == 1 + + +@pytest.mark.asyncio +async def test_terminate(fake_crate): + crate = await fake_crate(crate_handler()) + conn = await aio.connect(crate.url) + assert not inspect.iscoroutinefunction(conn.terminate) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + conn.terminate() + assert conn.closed + assert conn.client.closed + del conn + gc.collect() + await asyncio.sleep(0) + assert [str(w.message) for w in caught] == [] + + +@pytest.mark.asyncio +async def test_unclosed_connection_warns_once(fake_crate): + crate = await fake_crate(crate_handler()) + conn = await aio.connect(crate.url) + client = conn.client + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + del conn + gc.collect() + await asyncio.sleep(0) + assert [w.category for w in caught] == [ResourceWarning] + assert "unclosed