From ac38be0095700e87f0f6d9f80e9630c6632049ed Mon Sep 17 00:00:00 2001 From: Bilal Tonga Date: Thu, 15 Oct 2026 10:23:53 +0200 Subject: [PATCH 1/7] Move JSON encoding of the HTTP client into a shared module --- src/crate/client/_codec.py | 79 ++++++++++++++++++++++++++++++++++++++ src/crate/client/http.py | 62 +++--------------------------- 2 files changed, 85 insertions(+), 56 deletions(-) create mode 100644 src/crate/client/_codec.py diff --git a/src/crate/client/_codec.py b/src/crate/client/_codec.py new file mode 100644 index 00000000..ef07ea02 --- /dev/null +++ b/src/crate/client/_codec.py @@ -0,0 +1,79 @@ +# -*- 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 calendar +import datetime as dt +import typing as t +from decimal import Decimal + +import orjson + +epoch_aware = dt.datetime(1970, 1, 1, tzinfo=dt.timezone.utc) +epoch_naive = dt.datetime(1970, 1, 1) + + +def json_encoder(obj: t.Any) -> t.Union[int, str]: + """ + Encoder function for orjson, with additional type support. + + - Python's `Decimal` type will be serialized to `str`. + - Python's `dt.datetime` and `dt.date` types will be + serialized to `int` after converting to milliseconds + since epoch. + - Python's `dt.time` will be serialized to `str`, following + the ISO format. + + https://github.com/ijl/orjson#default + https://cratedb.com/docs/crate/reference/en/latest/general/ddl/data-types.html#type-timestamp + """ + if isinstance(obj, Decimal): + return str(obj) + if isinstance(obj, dt.datetime): + if obj.tzinfo is not None: + delta = obj - epoch_aware + else: + delta = obj - epoch_naive + return int( + delta.microseconds / 1000.0 + + (delta.seconds + delta.days * 24 * 3600) * 1000.0 + ) + if isinstance(obj, dt.time): + return obj.isoformat() + if isinstance(obj, dt.date): + return calendar.timegm(obj.timetuple()) * 1000 + raise TypeError + + +def json_dumps(obj: t.Any) -> bytes: + """ + Serialize to JSON format, using `orjson`, with additional type support. + + https://github.com/ijl/orjson + """ + return orjson.dumps( + obj, + default=json_encoder, + option=( + orjson.OPT_PASSTHROUGH_DATETIME + | orjson.OPT_NON_STR_KEYS + | orjson.OPT_SERIALIZE_NUMPY + ), + ) diff --git a/src/crate/client/http.py b/src/crate/client/http.py index 7492ecf2..5c877641 100644 --- a/src/crate/client/http.py +++ b/src/crate/client/http.py @@ -20,8 +20,6 @@ # software solely pursuant to the terms of the relevant commercial agreement. -import calendar -import datetime as dt import gzip import heapq import io @@ -33,7 +31,6 @@ import threading import typing as t from base64 import b64encode -from decimal import Decimal from time import time from urllib.parse import SplitResult, urlparse @@ -52,6 +49,12 @@ from urllib3.util.retry import Retry from verlib2 import Version +from crate.client._codec import ( # noqa: F401 + epoch_aware, + epoch_naive, + json_dumps, + json_encoder, +) from crate.client.exceptions import ( BlobLocationNotFoundException, ConnectionError, @@ -87,59 +90,6 @@ def super_len(o): return None -epoch_aware = dt.datetime(1970, 1, 1, tzinfo=dt.timezone.utc) -epoch_naive = dt.datetime(1970, 1, 1) - - -def json_encoder(obj: t.Any) -> t.Union[int, str]: - """ - Encoder function for orjson, with additional type support. - - - Python's `Decimal` type will be serialized to `str`. - - Python's `dt.datetime` and `dt.date` types will be - serialized to `int` after converting to milliseconds - since epoch. - - Python's `dt.time` will be serialized to `str`, following - the ISO format. - - https://github.com/ijl/orjson#default - https://cratedb.com/docs/crate/reference/en/latest/general/ddl/data-types.html#type-timestamp - """ - if isinstance(obj, Decimal): - return str(obj) - if isinstance(obj, dt.datetime): - if obj.tzinfo is not None: - delta = obj - epoch_aware - else: - delta = obj - epoch_naive - return int( - delta.microseconds / 1000.0 - + (delta.seconds + delta.days * 24 * 3600) * 1000.0 - ) - if isinstance(obj, dt.time): - return obj.isoformat() - if isinstance(obj, dt.date): - return calendar.timegm(obj.timetuple()) * 1000 - raise TypeError - - -def json_dumps(obj: t.Any) -> bytes: - """ - Serialize to JSON format, using `orjson`, with additional type support. - - https://github.com/ijl/orjson - """ - return orjson.dumps( - obj, - default=json_encoder, - option=( - orjson.OPT_PASSTHROUGH_DATETIME - | orjson.OPT_NON_STR_KEYS - | orjson.OPT_SERIALIZE_NUMPY - ), - ) - - class Server: def __init__(self, server, **pool_kw): socket_options = _get_socket_opts( From 054f4d109229e853e6ed9c5ffcdd034456c8efe7 Mon Sep 17 00:00:00 2001 From: Bilal Tonga Date: Thu, 15 Oct 2026 10:23:53 +0200 Subject: [PATCH 2/7] Move request building of the HTTP client into a shared module --- src/crate/client/_request.py | 131 +++++++++++++++++++++++++++++ src/crate/client/http.py | 96 +++++---------------- tests/client/test_serialization.py | 2 +- 3 files changed, 152 insertions(+), 77 deletions(-) create mode 100644 src/crate/client/_request.py diff --git a/src/crate/client/_request.py b/src/crate/client/_request.py new file mode 100644 index 00000000..975659f7 --- /dev/null +++ b/src/crate/client/_request.py @@ -0,0 +1,131 @@ +# -*- 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 gzip +import io +import os +import typing as t +from base64 import b64encode + +from crate.client._codec import json_dumps + + +def super_len(o: t.Any) -> t.Optional[int]: + if hasattr(o, "__len__"): + return len(o) + if hasattr(o, "len"): + return o.len + if hasattr(o, "fileno"): + try: + fileno = o.fileno() + except io.UnsupportedOperation: + pass + else: + return os.fstat(fileno).st_size + if hasattr(o, "getvalue"): + # e.g. BytesIO, cStringIO.StringI + return len(o.getvalue()) + return None + + +def build_headers( + headers: t.Optional[t.Dict[str, str]], + data: t.Any, + username: t.Optional[str] = None, + password: t.Optional[str] = None, + schema: t.Optional[str] = None, + jwt_token: t.Optional[str] = None, +) -> t.Dict[str, str]: + """ + Add the length, authentication, schema and content type headers. + + A given `headers` dict is updated in place and returned. + """ + if headers is None: + headers = {} + if "Content-Length" not in headers: + length = super_len(data) + if length is not None: + headers["Content-Length"] = str(length) + + # Sanity checks. + if jwt_token is not None and username is not None: + raise ValueError( + "Either JWT tokens are accepted, or user credentials, but not both" + ) + + # Authentication token + if jwt_token is not None and "Authorization" not in headers: + headers["Authorization"] = "Bearer %s" % jwt_token + + # Authentication credentials + if username is not None: + if "Authorization" not in headers and username is not None: + credentials = username + ":" + if password is not None: + credentials += password + headers["Authorization"] = "Basic %s" % b64encode( + credentials.encode("utf-8") + ).decode("utf-8") + # For backwards compatibility with Crate <= 2.2 + if "X-User" not in headers: + headers["X-User"] = username + + if schema is not None: + headers["Default-Schema"] = schema + headers["Accept"] = "application/json" + headers["Content-Type"] = "application/json" + return headers + + +def compress_body( + data: bytes, compress: t.Union[int, bool] +) -> t.Tuple[bytes, bool]: + """ + Return the request body to send and whether it is gzip-compressed. + + `compress` is either a flag, or the body size in bytes from which on + to compress. + """ + compress_enabled = compress is True or ( + not isinstance(compress, bool) and len(data) >= compress + ) + if compress_enabled: + return gzip.compress(data, compresslevel=6), True + return data, False + + +def create_sql_payload( + stmt: str, + args: t.Optional[t.Union[t.List, t.Tuple, t.Dict]], + bulk_args: t.Optional[t.List], +) -> bytes: + if not isinstance(stmt, str): + raise ValueError("stmt is not a string") + if args and bulk_args: + raise ValueError("Cannot provide both: args and bulk_args") + + data: t.Dict[str, t.Any] = {"stmt": stmt} + if args: + data["args"] = args + if bulk_args: + data["bulk_args"] = bulk_args + return json_dumps(data) diff --git a/src/crate/client/http.py b/src/crate/client/http.py index 5c877641..19049f40 100644 --- a/src/crate/client/http.py +++ b/src/crate/client/http.py @@ -20,9 +20,7 @@ # software solely pursuant to the terms of the relevant commercial agreement. -import gzip import heapq -import io import logging import os import re @@ -30,7 +28,6 @@ import ssl import threading import typing as t -from base64 import b64encode from time import time from urllib.parse import SplitResult, urlparse @@ -55,6 +52,15 @@ json_dumps, json_encoder, ) +from crate.client._request import ( # noqa: F401 + build_headers, + compress_body, + create_sql_payload, + super_len, +) +from crate.client._request import ( # noqa: F401 + create_sql_payload as _create_sql_payload, +) from crate.client.exceptions import ( BlobLocationNotFoundException, ConnectionError, @@ -72,24 +78,6 @@ SSL_ONLY_ARGS = {"ca_certs", "cert_reqs", "cert_file", "key_file"} -def super_len(o): - if hasattr(o, "__len__"): - return len(o) - if hasattr(o, "len"): - return o.len - if hasattr(o, "fileno"): - try: - fileno = o.fileno() - except io.UnsupportedOperation: - pass - else: - return os.fstat(fileno).st_size - if hasattr(o, "getvalue"): - # e.g. BytesIO, cStringIO.StringI - return len(o.getvalue()) - return None - - class Server: def __init__(self, server, **pool_kw): socket_options = _get_socket_opts( @@ -138,41 +126,14 @@ def request( path = "/{path_prefix}/{path}".format( path_prefix=self.path_prefix, path=path.strip("/") ) - if headers is None: - headers = {} - if "Content-Length" not in headers: - length = super_len(data) - if length is not None: - headers["Content-Length"] = str(length) - - # Sanity checks. - if jwt_token is not None and username is not None: - raise ValueError( - "Either JWT tokens are accepted, " - "or user credentials, but not both" - ) - - # Authentication token - if jwt_token is not None and "Authorization" not in headers: - headers["Authorization"] = "Bearer %s" % jwt_token - - # Authentication credentials - if username is not None: - if "Authorization" not in headers and username is not None: - credentials = username + ":" - if password is not None: - credentials += password - headers["Authorization"] = "Basic %s" % b64encode( - credentials.encode("utf-8") - ).decode("utf-8") - # For backwards compatibility with Crate <= 2.2 - if "X-User" not in headers: - headers["X-User"] = username - - if schema is not None: - headers["Default-Schema"] = schema - headers["Accept"] = "application/json" - headers["Content-Type"] = "application/json" + headers = build_headers( + headers, + data, + username=username, + password=password, + schema=schema, + jwt_token=jwt_token, + ) kwargs["assert_same_host"] = False kwargs["redirect"] = False kwargs["retries"] = Retry(read=0, backoff_factor=backoff_factor) @@ -344,20 +305,6 @@ def _update_pool_kwargs_for_ssl_minimum_version(server, kwargs): kwargs["ssl_minimum_version"] = ssl.TLSVersion.MINIMUM_SUPPORTED -def _create_sql_payload(stmt, args, bulk_args) -> bytes: - if not isinstance(stmt, str): - raise ValueError("stmt is not a string") - if args and bulk_args: - raise ValueError("Cannot provide both: args and bulk_args") - - data = {"stmt": stmt} - if args: - data["args"] = args - if bulk_args: - data["bulk_args"] = bulk_args - return json_dumps(data) - - def _get_socket_opts( keepalive=True, tcp_keepidle=None, tcp_keepintvl=None, tcp_keepcnt=None ): @@ -506,7 +453,7 @@ def sql(self, stmt, parameters=None, bulk_parameters=None): if stmt is None: return None - data = _create_sql_payload(stmt, parameters, bulk_parameters) + data = create_sql_payload(stmt, parameters, bulk_parameters) logger.debug("Sending request to %s with payload: %s", self.path, data) content = self._json_request("POST", self.path, data=data) logger.debug("JSON response for stmt(%s): %s", stmt, content) @@ -642,11 +589,8 @@ def _json_request(self, method, path, data): """ headers = {"Accept-Encoding": "gzip, deflate"} - compress_enabled = self.compress is True or ( - not isinstance(self.compress, bool) and len(data) >= self.compress - ) - if compress_enabled: - data = gzip.compress(data, compresslevel=6) + data, compressed = compress_body(data, self.compress) + if compressed: headers["Content-Encoding"] = "gzip" response = self._request(method, path, data=data, headers=headers) diff --git a/tests/client/test_serialization.py b/tests/client/test_serialization.py index ed917c4c..cc1243b9 100644 --- a/tests/client/test_serialization.py +++ b/tests/client/test_serialization.py @@ -46,7 +46,7 @@ def test_data_is_serialized(): """ mock = MagicMock(spec=bytes) - with patch("crate.client.http.json_dumps", return_value=mock) as f: + with patch("crate.client._request.json_dumps", return_value=mock) as f: with patch(REQUEST_PATH, return_value=fake_response(200)) as request: client = Client(servers="localhost:4200") client.sql( From 9c69b33106f8bd5e1fdbdb7b91a42811f3b5474d Mon Sep 17 00:00:00 2001 From: Bilal Tonga Date: Thu, 15 Oct 2026 10:23:54 +0200 Subject: [PATCH 3/7] Move response handling of the HTTP client into a shared module --- pyproject.toml | 2 +- src/crate/client/_response.py | 164 ++++++++++++++++++++++++++++++++++ src/crate/client/http.py | 116 ++++++++---------------- 3 files changed, 202 insertions(+), 80 deletions(-) create mode 100644 src/crate/client/_response.py diff --git a/pyproject.toml b/pyproject.toml index ed306181..4c46d752 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}" = [ +lint.per-file-ignores."src/crate/client/{connection.py,http.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/_response.py b/src/crate/client/_response.py new file mode 100644 index 00000000..b6dc6e95 --- /dev/null +++ b/src/crate/client/_response.py @@ -0,0 +1,164 @@ +# -*- 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 urllib.parse import urlparse + +import orjson + +from crate.client.exceptions import ( + BlobLocationNotFoundException, + ConnectionError, + IntegrityError, + ProgrammingError, +) + +SRV_UNAVAILABLE_STATUSES = {502, 503, 504, 509} +PRESERVE_ACTIVE_SERVER_EXCEPTIONS = {ConnectionResetError, BrokenPipeError} + + +def raise_for_status( + status: int, + reason: t.Optional[str], + headers: t.Mapping[str, str], + body: bytes, +) -> None: + """ + Raise `IntegrityError` exceptions for `DuplicateKeyException` errors. + """ + try: + return _raise_for_status_real(status, reason, headers, body) + except ProgrammingError as ex: + if "DuplicateKeyException" in ex.message: + raise IntegrityError(ex.message, error_trace=ex.error_trace) from ex + raise + + +def _raise_for_status_real( + status: int, + reason: t.Optional[str], + headers: t.Mapping[str, str], + body: bytes, +) -> None: + """make sure that only crate.exceptions are raised that are defined in + the DB-API specification""" + message = "" + if 400 <= status < 500: + message = "%s Client Error: %s" % (status, reason) + elif 500 <= status < 600: + message = "%s Server Error: %s" % (status, reason) + else: + return + if status == 503: + raise ConnectionError(message) + if headers.get("content-type", "").startswith("application/json"): + data = orjson.loads(body) + error = data.get("error", {}) + error_trace = data.get("error_trace", None) + if "results" in data: + errors = [ + res["error_message"] + for res in data["results"] + if res.get("error_message") + ] + if errors: + raise ProgrammingError("\n".join(errors)) + if isinstance(error, dict): + raise ProgrammingError( + error.get("message", ""), error_trace=error_trace + ) + raise ProgrammingError(error, error_trace=error_trace) + raise ProgrammingError(message) + + +def parse_server_infos( + server: str, content: t.Dict[str, t.Any] +) -> t.Tuple[str, t.Optional[str], str]: + node_name = content.get("name") + node_version = content.get("version", {}).get("number", "0.0.0") + return server, node_name, node_version + + +def blob_put_result( + status: int, + reason: t.Optional[str], + headers: t.Mapping[str, str], + body: bytes, + table: str, + digest: str, +) -> bool: + if status == 201: + # blob created + return True + if status == 409: + # blob exists + return False + if status in (400, 404): + raise BlobLocationNotFoundException(table, digest) + raise_for_status(status, reason, headers, body) + return False + + +def blob_del_result( + status: int, + reason: t.Optional[str], + headers: t.Mapping[str, str], + body: bytes, +) -> bool: + if status == 204: + return True + if status == 404: + return False + raise_for_status(status, reason, headers, body) + return False + + +def blob_exists_result( + status: int, + reason: t.Optional[str], + headers: t.Mapping[str, str], + body: bytes, +) -> bool: + if status == 200: + return True + elif status == 404: + return False + raise_for_status(status, reason, headers, body) + return False + + +def redirect_target( + status: int, location: t.Union[str, None, t.Literal[False]] +) -> t.Optional[str]: + """ + Return the server to follow a redirect to, as ``scheme://netloc``. + """ + if location and 300 <= status <= 308: + redirect_url = urlparse(location) + return f"{redirect_url.scheme}://{redirect_url.netloc}" + return None + + +def should_preserve_server(ex: BaseException) -> bool: + return any( + t in [type(arg) for arg in ex.args] + for t in PRESERVE_ACTIVE_SERVER_EXCEPTIONS + ) diff --git a/src/crate/client/http.py b/src/crate/client/http.py index 19049f40..e652ee7c 100644 --- a/src/crate/client/http.py +++ b/src/crate/client/http.py @@ -61,11 +61,20 @@ from crate.client._request import ( # noqa: F401 create_sql_payload as _create_sql_payload, ) +from crate.client._response import ( # noqa: F401 + PRESERVE_ACTIVE_SERVER_EXCEPTIONS, + SRV_UNAVAILABLE_STATUSES, + blob_del_result, + blob_exists_result, + blob_put_result, + parse_server_infos, + raise_for_status, + redirect_target, + should_preserve_server, +) from crate.client.exceptions import ( - BlobLocationNotFoundException, ConnectionError, DigestNotFoundException, - IntegrityError, ProgrammingError, ) @@ -73,8 +82,6 @@ _HTTP_PAT = pat = re.compile("https?://.+", re.I) -SRV_UNAVAILABLE_STATUSES = {502, 503, 504, 509} -PRESERVE_ACTIVE_SERVER_EXCEPTIONS = {ConnectionResetError, BrokenPipeError} SSL_ONLY_ARGS = {"ca_certs", "cert_reqs", "cert_file", "key_file"} @@ -171,47 +178,12 @@ def _ex_to_message(ex): def _raise_for_status(response): - """ - Raise `IntegrityError` exceptions for `DuplicateKeyException` errors. - """ - try: - return _raise_for_status_real(response) - except ProgrammingError as ex: - if "DuplicateKeyException" in ex.message: - raise IntegrityError(ex.message, error_trace=ex.error_trace) from ex - raise - - -def _raise_for_status_real(response): - """make sure that only crate.exceptions are raised that are defined in - the DB-API specification""" - message = "" - if 400 <= response.status < 500: - message = "%s Client Error: %s" % (response.status, response.reason) - elif 500 <= response.status < 600: - message = "%s Server Error: %s" % (response.status, response.reason) - else: - return - if response.status == 503: - raise ConnectionError(message) - if response.headers.get("content-type", "").startswith("application/json"): - data = orjson.loads(response.data) - error = data.get("error", {}) - error_trace = data.get("error_trace", None) - if "results" in data: - errors = [ - res["error_message"] - for res in data["results"] - if res.get("error_message") - ] - if errors: - raise ProgrammingError("\n".join(errors)) - if isinstance(error, dict): - raise ProgrammingError( - error.get("message", ""), error_trace=error_trace - ) - raise ProgrammingError(error, error_trace=error_trace) - raise ProgrammingError(message) + # Read the body of error responses only, so that a successful streamed + # response is not loaded into memory. + if 400 <= response.status < 600: + raise_for_status( + response.status, response.reason, response.headers, response.data + ) def _server_url(server): @@ -463,10 +435,7 @@ def sql(self, stmt, parameters=None, bulk_parameters=None): def server_infos(self, server): response = self._request("GET", "/", server=server) _raise_for_status(response) - content = _json_from_response(response) - node_name = content.get("name") - node_version = content.get("version", {}).get("number", "0.0.0") - return server, node_name, node_version + return parse_server_infos(server, _json_from_response(response)) def blob_put(self, table, digest, data) -> bool: """ @@ -474,28 +443,23 @@ def blob_put(self, table, digest, data) -> bool: given table and digest. """ response = self._request("PUT", _blob_path(table, digest), data=data) - if response.status == 201: - # blob created - return True - if response.status == 409: - # blob exists - return False - if response.status in (400, 404): - raise BlobLocationNotFoundException(table, digest) - _raise_for_status(response) - return False + return blob_put_result( + response.status, + response.reason, + response.headers, + response.data, + table, + digest, + ) def blob_del(self, table, digest) -> bool: """ Deletes the blob with given digest under the given table. """ response = self._request("DELETE", _blob_path(table, digest)) - if response.status == 204: - return True - if response.status == 404: - return False - _raise_for_status(response) - return False + return blob_del_result( + response.status, response.reason, response.headers, response.data + ) def blob_get(self, table, digest, chunk_size=1024 * 128): """ @@ -514,12 +478,9 @@ def blob_exists(self, table, digest) -> bool: under the given table. """ response = self._request("HEAD", _blob_path(table, digest)) - if response.status == 200: - return True - elif response.status == 404: - return False - _raise_for_status(response) - return False + return blob_exists_result( + response.status, response.reason, response.headers, response.data + ) def _add_server(self, server): with self._lock: @@ -544,10 +505,10 @@ def _request(self, method, path, server=None, **kwargs): jwt_token=self.jwt_token, **kwargs, ) - redirect_location = response.get_redirect_location() - if redirect_location and 300 <= response.status <= 308: - redirect_url = urlparse(redirect_location) - redirect_server = f"{redirect_url.scheme}://{redirect_url.netloc}" + redirect_server = redirect_target( + response.status, response.get_redirect_location() + ) + if redirect_server is not None: self._add_server(redirect_server) return self._request( method, path, server=redirect_server, **kwargs @@ -572,10 +533,7 @@ def _request(self, method, path, server=None, **kwargs): ) from ex preserve_server = False if isinstance(ex, ProtocolError): - preserve_server = any( - t in [type(arg) for arg in ex.args] - for t in PRESERVE_ACTIVE_SERVER_EXCEPTIONS - ) + preserve_server = should_preserve_server(ex) if not preserve_server: with self._lock: # drop server from active ones From c4d2a9251906823d167b641c5518d59a737a3055 Mon Sep 17 00:00:00 2001 From: Bilal Tonga Date: Thu, 15 Oct 2026 10:23:54 +0200 Subject: [PATCH 4/7] Move option parsing of the HTTP client into a shared module --- src/crate/client/_options.py | 120 +++++++++++++++++++++++++++++++++++ src/crate/client/http.py | 97 ++++++++-------------------- tests/client/test_http.py | 5 +- 3 files changed, 151 insertions(+), 71 deletions(-) create mode 100644 src/crate/client/_options.py diff --git a/src/crate/client/_options.py b/src/crate/client/_options.py new file mode 100644 index 00000000..6dcf8e10 --- /dev/null +++ b/src/crate/client/_options.py @@ -0,0 +1,120 @@ +# -*- 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 socket +import typing as t +from urllib.parse import urlparse + +SQL_PATH = "/_sql?types=true" + +_HTTP_PAT = re.compile("https?://.+", re.I) + + +def server_url(server: str) -> str: + """ + Normalizes a given server string to an url + + >>> print(server_url('a')) + http://a + >>> print(server_url('a/path')) + http://a/path + >>> print(server_url('a:9345')) + http://a:9345 + >>> print(server_url('a:9345/path')) + http://a:9345/path + >>> print(server_url('https://a:9345')) + https://a:9345 + >>> print(server_url('https://a:9345/path')) + https://a:9345/path + >>> print(server_url('https://a')) + https://a + >>> print(server_url('https://a/path')) + https://a/path + >>> print(server_url('demo.crate.io')) + http://demo.crate.io + """ + if not _HTTP_PAT.match(server): + server = "http://%s" % server + parsed = urlparse(server) + return parsed.geturl() + + +def to_server_list(servers: t.Union[str, t.Iterable[str]]) -> t.List[str]: + if isinstance(servers, str): + servers = servers.split() + return [server_url(s) for s in servers] + + +def credentials_from_url( + url: str, username: t.Optional[str], password: t.Optional[str] +) -> t.Tuple[t.Optional[str], t.Optional[str]]: + """ + Return the user name and password from the userinfo of `url`, falling + back to the given ones for the parts the URL does not contain. + """ + parsed = urlparse(url) + if parsed.username is not None: + username = parsed.username + if parsed.password is not None: + password = parsed.password + return username, password + + +def validate_compress(compress: t.Any) -> t.Union[int, bool]: + if not isinstance(compress, (bool, int)): + raise TypeError( + f"compress must be bool or int, got {type(compress).__name__!r}" + ) + return compress + + +def sql_path(error_trace: bool, path: str = SQL_PATH) -> str: + if error_trace: + path += "&error_trace=true" + return path + + +def keepalive_socket_options( + keepalive: bool = True, + tcp_keepidle: t.Optional[int] = None, + tcp_keepintvl: t.Optional[int] = None, + tcp_keepcnt: t.Optional[int] = None, +) -> t.List[t.Tuple[int, int, int]]: + """ + Return the `setsockopt` arguments for TCP keepalive, or an empty list + when `keepalive` is off. + """ + if not keepalive: + return [] + + # always use TCP keepalive + opts = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)] + + # hasattr check because some options depend on system capabilities + # see https://docs.python.org/3/library/socket.html#socket.SOMAXCONN + if hasattr(socket, "TCP_KEEPIDLE") and tcp_keepidle is not None: + opts.append((socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, tcp_keepidle)) + if hasattr(socket, "TCP_KEEPINTVL") and tcp_keepintvl is not None: + opts.append((socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, tcp_keepintvl)) + if hasattr(socket, "TCP_KEEPCNT") and tcp_keepcnt is not None: + opts.append((socket.IPPROTO_TCP, socket.TCP_KEEPCNT, tcp_keepcnt)) + return opts diff --git a/src/crate/client/http.py b/src/crate/client/http.py index e652ee7c..7509347e 100644 --- a/src/crate/client/http.py +++ b/src/crate/client/http.py @@ -23,8 +23,6 @@ import heapq import logging import os -import re -import socket import ssl import threading import typing as t @@ -52,15 +50,22 @@ json_dumps, json_encoder, ) +from crate.client._options import ( # noqa: F401 + _HTTP_PAT, + SQL_PATH, + credentials_from_url, + keepalive_socket_options, + server_url, + sql_path, + to_server_list, + validate_compress, +) from crate.client._request import ( # noqa: F401 build_headers, compress_body, create_sql_payload, super_len, ) -from crate.client._request import ( # noqa: F401 - create_sql_payload as _create_sql_payload, -) from crate.client._response import ( # noqa: F401 PRESERVE_ACTIVE_SERVER_EXCEPTIONS, SRV_UNAVAILABLE_STATUSES, @@ -80,8 +85,12 @@ logger = logging.getLogger(__name__) +# Private names that other code imports from this module. +_create_sql_payload = create_sql_payload +_server_url = server_url +_to_server_list = to_server_list +pat = _HTTP_PAT -_HTTP_PAT = pat = re.compile("https?://.+", re.I) SSL_ONLY_ARGS = {"ca_certs", "cert_reqs", "cert_file", "key_file"} @@ -186,41 +195,6 @@ def _raise_for_status(response): ) -def _server_url(server): - """ - Normalizes a given server string to an url - - >>> print(_server_url('a')) - http://a - >>> print(_server_url('a/path')) - http://a/path - >>> print(_server_url('a:9345')) - http://a:9345 - >>> print(_server_url('a:9345/path')) - http://a:9345/path - >>> print(_server_url('https://a:9345')) - https://a:9345 - >>> print(_server_url('https://a:9345/path')) - https://a:9345/path - >>> print(_server_url('https://a')) - https://a - >>> print(_server_url('https://a/path')) - https://a/path - >>> print(_server_url('demo.crate.io')) - http://demo.crate.io - """ - if not _HTTP_PAT.match(server): - server = "http://%s" % server - parsed = urlparse(server) - return parsed.geturl() - - -def _to_server_list(servers): - if isinstance(servers, str): - servers = servers.split() - return [_server_url(s) for s in servers] - - def _pool_kw_args( verify_ssl_cert, ca_cert, @@ -287,20 +261,12 @@ def _get_socket_opts( if not keepalive: return None - # always use TCP keepalive - opts = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)] - - # hasattr check because some options depend on system capabilities - # see https://docs.python.org/3/library/socket.html#socket.SOMAXCONN - if hasattr(socket, "TCP_KEEPIDLE") and tcp_keepidle is not None: - opts.append((socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, tcp_keepidle)) - if hasattr(socket, "TCP_KEEPINTVL") and tcp_keepintvl is not None: - opts.append((socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, tcp_keepintvl)) - if hasattr(socket, "TCP_KEEPCNT") and tcp_keepcnt is not None: - opts.append((socket.IPPROTO_TCP, socket.TCP_KEEPCNT, tcp_keepcnt)) - # additionally use urllib3's default socket options - return list(HTTPConnection.default_socket_options) + opts + return list(HTTPConnection.default_socket_options) + ( + keepalive_socket_options( + keepalive, tcp_keepidle, tcp_keepintvl, tcp_keepcnt + ) + ) class Client: @@ -308,7 +274,7 @@ class Client: Crate connection client using CrateDB's HTTP API. """ - SQL_PATH = "/_sql?types=true" + SQL_PATH = SQL_PATH """Crate URI path for issuing SQL statements.""" retry_interval = 30 @@ -342,18 +308,15 @@ def __init__( if not servers: servers = [self.default_server] else: - servers = _to_server_list(servers) + servers = to_server_list(servers) # Try to derive credentials from first server argument if not # explicitly given. if servers and not username: try: - url = urlparse(servers[0]) - - if url.username is not None: - username = url.username - if url.password is not None: - password = url.password + username, password = credentials_from_url( + servers[0], username, password + ) except Exception as ex: logger.warning( "Unable to decode credentials from database " @@ -391,15 +354,9 @@ def __init__( self.jwt_token = jwt_token self.schema = schema - if not isinstance(compress, (bool, int)): - raise TypeError( - f"compress must be bool or int, got {type(compress).__name__!r}" - ) - self.compress = compress + self.compress = validate_compress(compress) - self.path = self.SQL_PATH - if error_trace: - self.path += "&error_trace=true" + self.path = sql_path(error_trace, self.SQL_PATH) def close(self): for server in self.server_pool.values(): diff --git a/tests/client/test_http.py b/tests/client/test_http.py index 3510de41..2e5bce94 100644 --- a/tests/client/test_http.py +++ b/tests/client/test_http.py @@ -281,7 +281,10 @@ def test_credentials_derived(): assert client.username == expected_user assert client.password == expected_password - with patch("crate.client.http.urlparse", side_effect=Exception): + with ( + patch("crate.client._options.urlparse", side_effect=Exception), + patch("crate.client.http.urlparse", side_effect=Exception), + ): Client("") actual_username = "actual_username" From db2d13f02376cdaf1d3755705f4d7e1638191455 Mon Sep 17 00:00:00 2001 From: Bilal Tonga Date: Thu, 15 Oct 2026 10:23:54 +0200 Subject: [PATCH 5/7] Move server selection of the HTTP client into a shared pool class --- pyproject.toml | 2 +- src/crate/client/_pool.py | 120 ++++++++++++++++++++++++++++++++++++++ src/crate/client/http.py | 80 ++----------------------- 3 files changed, 126 insertions(+), 76 deletions(-) create mode 100644 src/crate/client/_pool.py diff --git a/pyproject.toml b/pyproject.toml index 4c46d752..13bacc48 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,_response.py}" = [ +lint.per-file-ignores."src/crate/client/{connection.py,http.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/_pool.py b/src/crate/client/_pool.py new file mode 100644 index 00000000..ac7489dc --- /dev/null +++ b/src/crate/client/_pool.py @@ -0,0 +1,120 @@ +# -*- 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 heapq +import logging +import threading +import time +import typing as t + +from crate.client.exceptions import ConnectionError + + +class _ServerPool: + """ + Round-robin server selection. A dropped server is taken out of rotation + for `retry_interval` seconds. + + Subclasses set `_logger` to log server changes on their own logger. + """ + + retry_interval = 30 + """Retry interval for failed servers in seconds.""" + + _logger = logging.getLogger(__name__) + + def __init__( + self, + servers: t.List[str], + clock: t.Callable[[], float] = time.time, + ) -> None: + self._active_servers = servers + self._inactive_servers: t.List[t.Tuple[float, str, str]] = [] + self._lock = threading.RLock() + self._clock = clock + + def _get_server(self) -> str: + """ + Get server to use for request. + Also process inactive server list, re-add them after given interval. + """ + with self._lock: + inactive_server_count = len(self._inactive_servers) + for _ in range(inactive_server_count): + try: + ts, server, message = heapq.heappop(self._inactive_servers) + except IndexError: + pass + else: + if (ts + self.retry_interval) > self._clock(): + # Not yet, put it back + heapq.heappush( + self._inactive_servers, (ts, server, message) + ) + else: + self._active_servers.append(server) + self._logger.warning( + "Restored server %s into active pool", server + ) + + # if none is old enough, use oldest + if not self._active_servers: + ts, server, message = heapq.heappop(self._inactive_servers) + self._active_servers.append(server) + self._logger.info("Restored server %s into active pool", server) + + server = self._active_servers[0] + self._roundrobin() + + return server + + @property + def active_servers(self) -> t.List[str]: + """get the active servers for this client""" + with self._lock: + return list(self._active_servers) + + def _drop_server(self, server: str, message: str) -> None: + """ + Drop server from active list and adds it to the inactive ones. + """ + try: + self._active_servers.remove(server) + except ValueError: + pass + else: + heapq.heappush( + self._inactive_servers, (self._clock(), server, message) + ) + self._logger.warning("Removed server %s from active pool", server) + + # if this is the last server raise exception, otherwise try next + if not self._active_servers: + raise ConnectionError( + ("No more Servers available, exception from last server: %s") + % message + ) + + def _roundrobin(self) -> None: + """ + Very simple round-robin implementation + """ + self._active_servers.append(self._active_servers.pop(0)) diff --git a/src/crate/client/http.py b/src/crate/client/http.py index 7509347e..39697166 100644 --- a/src/crate/client/http.py +++ b/src/crate/client/http.py @@ -20,13 +20,11 @@ # software solely pursuant to the terms of the relevant commercial agreement. -import heapq import logging import os import ssl import threading import typing as t -from time import time from urllib.parse import SplitResult, urlparse import orjson @@ -60,6 +58,7 @@ to_server_list, validate_compress, ) +from crate.client._pool import _ServerPool from crate.client._request import ( # noqa: F401 build_headers, compress_body, @@ -269,7 +268,7 @@ def _get_socket_opts( ) -class Client: +class Client(_ServerPool): """ Crate connection client using CrateDB's HTTP API. """ @@ -277,12 +276,11 @@ class Client: SQL_PATH = SQL_PATH """Crate URI path for issuing SQL statements.""" - retry_interval = 30 - """Retry interval for failed servers in seconds.""" - default_server = "http://127.0.0.1:4200" """Default server to use if no servers are given on instantiation.""" + _logger = logger + def __init__( self, servers=None, @@ -324,8 +322,7 @@ def __init__( "authentication: {ex}".format(ex=ex) ) - self._active_servers = servers - self._inactive_servers: t.List[t.Tuple[float, str, str]] = [] + super().__init__(servers) pool_kw = _pool_kw_args( verify_ssl_cert, ca_cert, @@ -347,7 +344,6 @@ def __init__( self.server_pool: t.Dict[str, Server] = {} self._update_server_pool(servers, **pool_kw) self._pool_kw = pool_kw - self._lock = threading.RLock() self._local = threading.local() self.username = username self.password = password @@ -514,71 +510,5 @@ def _json_request(self, method, path, data): return _json_from_response(response) return response.data - def _get_server(self): - """ - Get server to use for request. - Also process inactive server list, re-add them after given interval. - """ - with self._lock: - inactive_server_count = len(self._inactive_servers) - for _ in range(inactive_server_count): - try: - ts, server, message = heapq.heappop(self._inactive_servers) - except IndexError: - pass - else: - if (ts + self.retry_interval) > time(): - # Not yet, put it back - heapq.heappush( - self._inactive_servers, (ts, server, message) - ) - else: - self._active_servers.append(server) - logger.warning( - "Restored server %s into active pool", server - ) - - # if none is old enough, use oldest - if not self._active_servers: - ts, server, message = heapq.heappop(self._inactive_servers) - self._active_servers.append(server) - logger.info("Restored server %s into active pool", server) - - server = self._active_servers[0] - self._roundrobin() - - return server - - @property - def active_servers(self): - """get the active servers for this client""" - with self._lock: - return list(self._active_servers) - - def _drop_server(self, server, message): - """ - Drop server from active list and adds it to the inactive ones. - """ - try: - self._active_servers.remove(server) - except ValueError: - pass - else: - heapq.heappush(self._inactive_servers, (time(), server, message)) - logger.warning("Removed server %s from active pool", server) - - # if this is the last server raise exception, otherwise try next - if not self._active_servers: - raise ConnectionError( - ("No more Servers available, exception from last server: %s") - % message - ) - - def _roundrobin(self): - """ - Very simple round-robin implementation - """ - self._active_servers.append(self._active_servers.pop(0)) - def __repr__(self): return "".format(str(self._active_servers)) From 417a91d36ea39ef79517aefebdbe75839beb68fd Mon Sep 17 00:00:00 2001 From: Bilal Tonga Date: Thu, 15 Oct 2026 10:23:55 +0200 Subject: [PATCH 6/7] Add tests for the shared HTTP client modules --- tests/client/test_options.py | 149 +++++++++++++++++++++++ tests/client/test_pool.py | 148 +++++++++++++++++++++++ tests/client/test_request.py | 177 +++++++++++++++++++++++++++ tests/client/test_response.py | 217 ++++++++++++++++++++++++++++++++++ 4 files changed, 691 insertions(+) create mode 100644 tests/client/test_options.py create mode 100644 tests/client/test_pool.py create mode 100644 tests/client/test_request.py create mode 100644 tests/client/test_response.py diff --git a/tests/client/test_options.py b/tests/client/test_options.py new file mode 100644 index 00000000..56c39a85 --- /dev/null +++ b/tests/client/test_options.py @@ -0,0 +1,149 @@ +# -*- 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 socket + +import pytest + +from crate.client._options import ( + credentials_from_url, + keepalive_socket_options, + server_url, + sql_path, + to_server_list, + validate_compress, +) + + +@pytest.mark.parametrize( + "server, expected", + [ + ("a", "http://a"), + ("a/path", "http://a/path"), + ("a:9345", "http://a:9345"), + ("a:9345/path", "http://a:9345/path"), + ("https://a:9345", "https://a:9345"), + ("https://a:9345/path", "https://a:9345/path"), + ("https://a", "https://a"), + ("https://a/path", "https://a/path"), + ("HTTPS://a", "https://a"), + ("demo.crate.io", "http://demo.crate.io"), + ], +) +def test_server_url(server, expected): + assert server_url(server) == expected + + +@pytest.mark.parametrize( + "servers", + [ + "localhost:4200 https://localhost:4201", + ["localhost:4200", "https://localhost:4201"], + ("localhost:4200", "https://localhost:4201"), + ], +) +def test_to_server_list(servers): + assert to_server_list(servers) == [ + "http://localhost:4200", + "https://localhost:4201", + ] + + +@pytest.mark.parametrize( + "url, username, password, expected", + [ + ("http://localhost:4200", None, None, (None, None)), + ("http://crate@localhost:4200", None, None, ("crate", None)), + ( + "http://crate:s3cr3t@localhost:4200", + None, + None, + ("crate", "s3cr3t"), + ), + ("http://localhost:4200", None, "s3cr3t", (None, "s3cr3t")), + ("http://crate@localhost:4200", None, "s3cr3t", ("crate", "s3cr3t")), + ( + "http://crate:from-url@localhost:4200", + None, + "s3cr3t", + ("crate", "from-url"), + ), + ("http://localhost:4200", "", None, ("", None)), + ], +) +def test_credentials_from_url(url, username, password, expected): + assert credentials_from_url(url, username, password) == expected + + +@pytest.mark.parametrize("compress", [True, False, 0, 8192]) +def test_validate_compress(compress): + assert validate_compress(compress) is compress + + +@pytest.mark.parametrize( + "compress, type_name", [("8192", "'str'"), (None, "'NoneType'")] +) +def test_validate_compress_invalid(compress, type_name): + with pytest.raises( + TypeError, match=f"compress must be bool or int, got {type_name}" + ): + validate_compress(compress) + + +@pytest.mark.parametrize( + "error_trace, expected", + [ + (False, "/_sql?types=true"), + (True, "/_sql?types=true&error_trace=true"), + ], +) +def test_sql_path(error_trace, expected): + assert sql_path(error_trace) == expected + + +def test_sql_path_with_given_path(): + assert sql_path(True, "/_sql") == "/_sql&error_trace=true" + + +def test_keepalive_socket_options_off(): + assert keepalive_socket_options(False, 30, 10, 5) == [] + + +def test_keepalive_socket_options_defaults(): + assert keepalive_socket_options() == [ + (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) + ] + + +@pytest.mark.skipif( + not all( + hasattr(socket, name) + for name in ("TCP_KEEPIDLE", "TCP_KEEPINTVL", "TCP_KEEPCNT") + ), + reason="Platform does not support all TCP keepalive options", +) +def test_keepalive_socket_options_tcp(): + assert keepalive_socket_options(True, 30, 10, 5) == [ + (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1), + (socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 30), + (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 10), + (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 5), + ] diff --git a/tests/client/test_pool.py b/tests/client/test_pool.py new file mode 100644 index 00000000..7887e273 --- /dev/null +++ b/tests/client/test_pool.py @@ -0,0 +1,148 @@ +# -*- 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 logging + +import pytest + +from crate.client import exceptions +from crate.client._pool import _ServerPool +from crate.client.http import Client + +NODE1 = "http://localhost:4200" +NODE2 = "http://localhost:4201" +NODE3 = "http://localhost:4202" + + +class FakeClock: + def __init__(self, now: float = 1_700_000_000.0): + self.now = now + + def __call__(self) -> float: + return self.now + + +@pytest.fixture +def clock(): + return FakeClock() + + +def test_round_robin(clock): + pool = _ServerPool([NODE1, NODE2, NODE3], clock=clock) + picked = [pool._get_server() for _ in range(4)] + assert picked == [NODE1, NODE2, NODE3, NODE1] + + +def test_active_servers_is_a_copy(clock): + pool = _ServerPool([NODE1, NODE2], clock=clock) + active = pool.active_servers + active.clear() + assert pool.active_servers == [NODE1, NODE2] + + +def test_drop_server(clock): + pool = _ServerPool([NODE1, NODE2], clock=clock) + pool._drop_server(NODE1, "Connection refused") + assert pool.active_servers == [NODE2] + assert pool._inactive_servers == [(clock.now, NODE1, "Connection refused")] + assert [pool._get_server() for _ in range(2)] == [NODE2, NODE2] + + +def test_drop_unknown_server_is_ignored(clock): + pool = _ServerPool([NODE1, NODE2], clock=clock) + pool._drop_server(NODE3, "Connection refused") + assert pool.active_servers == [NODE1, NODE2] + assert pool._inactive_servers == [] + + +def test_drop_last_server_raises(clock): + pool = _ServerPool([NODE1, NODE2], clock=clock) + pool._drop_server(NODE1, "Connection refused") + with pytest.raises( + exceptions.ConnectionError, + match="No more Servers available, exception from last server: " + "Read timed out", + ): + pool._drop_server(NODE2, "Read timed out") + assert pool.active_servers == [] + + +def test_server_restored_after_retry_interval(clock): + pool = _ServerPool([NODE1, NODE2], clock=clock) + pool.retry_interval = 5 + pool._drop_server(NODE1, "Connection refused") + + clock.now += 4 + pool._get_server() + assert pool.active_servers == [NODE2] + + clock.now += 1 + assert pool._get_server() == NODE2 + assert pool.active_servers == [NODE1, NODE2] + assert pool._inactive_servers == [] + + +def test_retry_interval_class_default(clock): + pool = _ServerPool([NODE1, NODE2], clock=clock) + pool._drop_server(NODE1, "Connection refused") + + clock.now += 29 + pool._get_server() + assert pool.active_servers == [NODE2] + + clock.now += 1 + pool._get_server() + assert sorted(pool.active_servers) == [NODE1, NODE2] + + +def test_oldest_server_restored_when_none_is_active(clock): + pool = _ServerPool([NODE1, NODE2, NODE3], clock=clock) + pool._drop_server(NODE2, "Connection refused") + clock.now += 1 + pool._drop_server(NODE1, "Connection refused") + pool._active_servers = [] + + assert pool._get_server() == NODE2 + assert pool.active_servers == [NODE2] + assert pool._inactive_servers == [(clock.now, NODE1, "Connection refused")] + + +def test_pool_logs_on_subclass_logger(clock, caplog): + class Pool(_ServerPool): + _logger = logging.getLogger("crate.client.test") + + pool = Pool([NODE1, NODE2], clock=clock) + pool.retry_interval = 5 + with caplog.at_level(logging.WARNING): + pool._drop_server(NODE1, "Connection refused") + clock.now += 5 + pool._get_server() + assert [(r.name, r.getMessage()) for r in caplog.records] == [ + ("crate.client.test", f"Removed server {NODE1} from active pool"), + ("crate.client.test", f"Restored server {NODE1} into active pool"), + ] + + +def test_client_logs_on_http_logger(caplog): + client = Client(servers=[NODE1, NODE2]) + with caplog.at_level(logging.WARNING): + client._drop_server(NODE1, "Connection refused") + assert [r.name for r in caplog.records] == ["crate.client.http"] diff --git a/tests/client/test_request.py b/tests/client/test_request.py new file mode 100644 index 00000000..0423a4f2 --- /dev/null +++ b/tests/client/test_request.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 gzip + +import orjson +import pytest + +from crate.client._request import ( + build_headers, + compress_body, + create_sql_payload, +) + +STMT = b'{"stmt":"SELECT name FROM sys.cluster"}' +JWT = "eyJhbGciOiJSUzI1NiJ9.eyJpc3MiOiJjcmF0ZWRiIn0.c2lnbmF0dXJl" + + +def test_build_headers_without_credentials(): + assert build_headers(None, STMT) == { + "Content-Length": str(len(STMT)), + "Accept": "application/json", + "Content-Type": "application/json", + } + + +def test_build_headers_keeps_content_length(): + headers = build_headers({"Content-Length": "1024"}, STMT) + assert headers["Content-Length"] == "1024" + + +def test_build_headers_without_body_length(): + assert "Content-Length" not in build_headers(None, None) + + +def test_build_headers_updates_given_dict(): + headers = {"Accept-Encoding": "gzip, deflate"} + assert build_headers(headers, STMT) is headers + assert headers["Accept-Encoding"] == "gzip, deflate" + + +@pytest.mark.parametrize( + "password, authorization", + [ + # base64 of "crate:" + (None, "Basic Y3JhdGU6"), + # base64 of "crate:s3cr3t" + ("s3cr3t", "Basic Y3JhdGU6czNjcjN0"), + ], +) +def test_build_headers_basic_auth(password, authorization): + headers = build_headers(None, STMT, username="crate", password=password) + assert headers["Authorization"] == authorization + assert headers["X-User"] == "crate" + + +def test_build_headers_password_without_username(): + headers = build_headers(None, STMT, password="s3cr3t") + assert "Authorization" not in headers + assert "X-User" not in headers + + +def test_build_headers_jwt(): + headers = build_headers(None, STMT, jwt_token=JWT) + assert headers["Authorization"] == f"Bearer {JWT}" + assert "X-User" not in headers + + +def test_build_headers_jwt_and_username(): + with pytest.raises(ValueError, match="Either JWT tokens are accepted"): + build_headers(None, STMT, username="crate", jwt_token=JWT) + + +@pytest.mark.parametrize( + "credentials", + [ + {"username": "crate", "password": "s3cr3t"}, + {"jwt_token": JWT}, + ], +) +def test_build_headers_existing_authorization_wins(credentials): + headers = build_headers( + {"Authorization": "Bearer from-caller"}, STMT, **credentials + ) + assert headers["Authorization"] == "Bearer from-caller" + + +def test_build_headers_existing_x_user_wins(): + headers = build_headers({"X-User": "admin"}, STMT, username="crate") + assert headers["X-User"] == "admin" + + +def test_build_headers_default_schema(): + headers = build_headers(None, STMT, schema="doc") + assert headers["Default-Schema"] == "doc" + + +def test_build_headers_overrides_content_type(): + headers = build_headers( + {"Accept": "text/plain", "Content-Type": "text/plain"}, STMT + ) + assert headers["Accept"] == "application/json" + assert headers["Content-Type"] == "application/json" + + +@pytest.mark.parametrize( + "compress, compressed", + [ + (True, True), + (False, False), + (len(STMT), True), + (len(STMT) + 1, False), + (0, True), + ], +) +def test_compress_body(compress, compressed): + data, is_compressed = compress_body(STMT, compress) + assert is_compressed is compressed + if compressed: + assert gzip.decompress(data) == STMT + else: + assert data is STMT + + +@pytest.mark.parametrize( + "args, bulk_args, expected", + [ + (None, None, {"stmt": "INSERT INTO doc.t (x) VALUES (?)"}), + ( + [1], + None, + {"stmt": "INSERT INTO doc.t (x) VALUES (?)", "args": [1]}, + ), + ( + None, + [[1], [2]], + { + "stmt": "INSERT INTO doc.t (x) VALUES (?)", + "bulk_args": [[1], [2]], + }, + ), + ([], [], {"stmt": "INSERT INTO doc.t (x) VALUES (?)"}), + ], +) +def test_create_sql_payload(args, bulk_args, expected): + payload = create_sql_payload( + "INSERT INTO doc.t (x) VALUES (?)", args, bulk_args + ) + assert orjson.loads(payload) == expected + + +def test_create_sql_payload_stmt_not_a_string(): + with pytest.raises(ValueError, match="stmt is not a string"): + create_sql_payload(b"SELECT 1", None, None) + + +def test_create_sql_payload_args_and_bulk_args(): + with pytest.raises(ValueError, match="Cannot provide both"): + create_sql_payload("INSERT INTO doc.t (x) VALUES (?)", [1], [[1], [2]]) diff --git a/tests/client/test_response.py b/tests/client/test_response.py new file mode 100644 index 00000000..cfbf7927 --- /dev/null +++ b/tests/client/test_response.py @@ -0,0 +1,217 @@ +# -*- 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 orjson +import pytest + +from crate.client import exceptions +from crate.client._response import ( + blob_del_result, + blob_exists_result, + blob_put_result, + parse_server_infos, + raise_for_status, + redirect_target, + should_preserve_server, +) + +JSON = {"content-type": "application/json"} +TEXT = {"content-type": "text/plain"} +DIGEST = "c520e6109835c876fd98636efec43dd61634b7d3" + + +def sql_error(message, code=4000, error_trace=None): + data = {"error": {"message": message, "code": code}} + if error_trace is not None: + data["error_trace"] = error_trace + return orjson.dumps(data) + + +@pytest.mark.parametrize("status", [100, 200, 201, 204, 301, 307, 600]) +def test_raise_for_status_no_error(status): + assert raise_for_status(status, "", JSON, b"") is None + + +@pytest.mark.parametrize( + "status, reason, message", + [ + (400, "Bad Request", "400 Client Error: Bad Request"), + (404, "Not Found", "404 Client Error: Not Found"), + (500, "Internal Server Error", "500 Server Error: Internal Server"), + (502, "Bad Gateway", "502 Server Error: Bad Gateway"), + ], +) +def test_raise_for_status_without_json_body(status, reason, message): + with pytest.raises(exceptions.ProgrammingError, match=message): + raise_for_status(status, reason, TEXT, b"") + + +@pytest.mark.parametrize("headers", [JSON, TEXT]) +def test_raise_for_status_503(headers): + with pytest.raises( + exceptions.ConnectionError, + match="503 Server Error: Service Unavailable", + ): + raise_for_status(503, "Service Unavailable", headers, b"") + + +def test_raise_for_status_sql_error(): + body = sql_error( + "SQLParseException[line 1:1: mismatched input 'SELCT']", + error_trace="io.crate.exceptions.SQLParseException: ...", + ) + with pytest.raises(exceptions.ProgrammingError) as excinfo: + raise_for_status(400, "Bad Request", JSON, body) + assert type(excinfo.value) is exceptions.ProgrammingError + assert excinfo.value.message == ( + "SQLParseException[line 1:1: mismatched input 'SELCT']" + ) + assert excinfo.value.error_trace == ( + "io.crate.exceptions.SQLParseException: ..." + ) + + +def test_raise_for_status_error_as_string(): + body = orjson.dumps({"error": "No handler found for uri [/_sq]"}) + with pytest.raises( + exceptions.ProgrammingError, match=r"No handler found for uri \[/_sq\]" + ): + raise_for_status(400, "Bad Request", JSON, body) + + +def test_raise_for_status_duplicate_key(): + body = sql_error( + "DuplicateKeyException[A document with the same primary key " + "exists already]", + code=4091, + error_trace="trace", + ) + with pytest.raises(exceptions.IntegrityError) as excinfo: + raise_for_status(409, "Conflict", JSON, body) + assert excinfo.value.message.startswith("DuplicateKeyException[") + assert excinfo.value.error_trace == "trace" + + +def test_raise_for_status_bulk_errors_are_joined(): + body = orjson.dumps( + { + "results": [ + {"rowcount": 1}, + {"rowcount": -2, "error_message": "first failed"}, + {"rowcount": -2, "error_message": "second failed"}, + ] + } + ) + with pytest.raises(exceptions.ProgrammingError) as excinfo: + raise_for_status(400, "Bad Request", JSON, body) + assert excinfo.value.message == "first failed\nsecond failed" + + +def test_raise_for_status_bulk_without_errors_uses_error(): + body = orjson.dumps( + {"results": [{"rowcount": 1}], "error": {"message": "failed"}} + ) + with pytest.raises(exceptions.ProgrammingError, match="^failed$"): + raise_for_status(400, "Bad Request", JSON, body) + + +@pytest.mark.parametrize( + "content, expected", + [ + ( + {"name": "crate-1", "version": {"number": "6.1.2"}}, + ("http://localhost:4200", "crate-1", "6.1.2"), + ), + ({}, ("http://localhost:4200", None, "0.0.0")), + ], +) +def test_parse_server_infos(content, expected): + assert parse_server_infos("http://localhost:4200", content) == expected + + +@pytest.mark.parametrize( + "status, expected", [(201, True), (409, False), (200, False)] +) +def test_blob_put_result(status, expected): + assert blob_put_result(status, "", TEXT, b"", "myfiles", DIGEST) is ( + expected + ) + + +@pytest.mark.parametrize("status", [400, 404]) +def test_blob_put_result_location_not_found(status): + with pytest.raises(exceptions.BlobLocationNotFoundException): + blob_put_result(status, "", TEXT, b"", "myfiles", DIGEST) + + +@pytest.mark.parametrize( + "status, expected", [(204, True), (404, False), (200, False)] +) +def test_blob_del_result(status, expected): + assert blob_del_result(status, "", TEXT, b"") is expected + + +@pytest.mark.parametrize( + "status, expected", [(200, True), (404, False), (204, False)] +) +def test_blob_exists_result(status, expected): + assert blob_exists_result(status, "", TEXT, b"") is expected + + +@pytest.mark.parametrize( + "result", [blob_put_result, blob_del_result, blob_exists_result] +) +def test_blob_result_server_error(result): + args = (500, "Internal Server Error", TEXT, b"") + if result is blob_put_result: + args += ("myfiles", DIGEST) + with pytest.raises(exceptions.ProgrammingError, match="500 Server Error"): + result(*args) + + +@pytest.mark.parametrize( + "status, location, expected", + [ + (307, "http://localhost:4201/_sql", "http://localhost:4201"), + (301, "https://node2:4200/_blobs/x", "https://node2:4200"), + (308, "http://localhost:4201/", "http://localhost:4201"), + (307, None, None), + (307, "", None), + (200, False, None), + (309, "http://localhost:4201/_sql", None), + ], +) +def test_redirect_target(status, location, expected): + assert redirect_target(status, location) == expected + + +@pytest.mark.parametrize( + "ex, expected", + [ + (Exception("reset", ConnectionResetError(104, "reset")), True), + (Exception("broken", BrokenPipeError(32, "broken")), True), + (Exception("refused", ConnectionRefusedError()), False), + (Exception("Connection aborted."), False), + (ConnectionResetError(104, "reset"), False), + ], +) +def test_should_preserve_server(ex, expected): + assert should_preserve_server(ex) is expected From 0610d282bdd33444d47f275420fc370b59baa7a7 Mon Sep 17 00:00:00 2001 From: Bilal Tonga Date: Thu, 15 Oct 2026 13:51:59 +0200 Subject: [PATCH 7/7] Drop unused private aliases from the HTTP client module --- src/crate/client/http.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/src/crate/client/http.py b/src/crate/client/http.py index 39697166..4259628f 100644 --- a/src/crate/client/http.py +++ b/src/crate/client/http.py @@ -48,12 +48,10 @@ json_dumps, json_encoder, ) -from crate.client._options import ( # noqa: F401 - _HTTP_PAT, +from crate.client._options import ( SQL_PATH, credentials_from_url, keepalive_socket_options, - server_url, sql_path, to_server_list, validate_compress, @@ -84,12 +82,6 @@ logger = logging.getLogger(__name__) -# Private names that other code imports from this module. -_create_sql_payload = create_sql_payload -_server_url = server_url -_to_server_list = to_server_list -pat = _HTTP_PAT - SSL_ONLY_ARGS = {"ca_certs", "cert_reqs", "cert_file", "key_file"}