diff options
Diffstat (limited to 'contrib/python')
48 files changed, 1790 insertions, 491 deletions
diff --git a/contrib/python/Flask-Cors/py3/.dist-info/METADATA b/contrib/python/Flask-Cors/py3/.dist-info/METADATA index 513e909c173..39686d45aa6 100644 --- a/contrib/python/Flask-Cors/py3/.dist-info/METADATA +++ b/contrib/python/Flask-Cors/py3/.dist-info/METADATA @@ -1,6 +1,6 @@ Metadata-Version: 2.4 Name: flask-cors -Version: 6.0.3 +Version: 6.0.5 Summary: A Flask extension simplifying CORS support Author-email: Cory Dolphin <[email protected]> License-Expression: MIT @@ -21,6 +21,7 @@ Requires-Python: <4.0,>=3.9 Description-Content-Type: text/x-rst Requires-Dist: flask>=0.9 Requires-Dist: Werkzeug>=0.7 +Requires-Dist: typing_extensions>=4.6.0; python_version < "3.11" Flask-CORS ========== diff --git a/contrib/python/Flask-Cors/py3/flask_cors/core.py b/contrib/python/Flask-Cors/py3/flask_cors/core.py index 8df343c935a..1eeb20a83a6 100644 --- a/contrib/python/Flask-Cors/py3/flask_cors/core.py +++ b/contrib/python/Flask-Cors/py3/flask_cors/core.py @@ -1,13 +1,106 @@ +from __future__ import annotations + import logging import re -from collections.abc import Iterable +from collections.abc import Iterable, Mapping +from dataclasses import dataclass from datetime import timedelta +from typing import Any, TypedDict, Union, cast -from flask import current_app, request +from flask import Blueprint, Flask, Response, current_app, request from werkzeug.datastructures import Headers, MultiDict LOG = logging.getLogger(__name__) +# A resource/origin/header pattern may either be a literal string or a +# pre-compiled regular expression. +ResourcePattern = Union[str, "re.Pattern[str]"] + +# What a user may pass as ``resources``: a single pattern, a list of patterns, +# or a mapping of pattern -> per-resource (raw) options. +ResourceSpec = Union[ + str, + "re.Pattern[str]", + "list[ResourcePattern]", + "dict[ResourcePattern, dict[str, Any]]", +] + +# Input shapes accepted from users for individual options, before they are +# normalized by :func:`serialize_options`. These are deliberately broad: a +# scalar or an iterable of scalars are both accepted in most cases. +OriginsInput = Union[ResourcePattern, "Iterable[ResourcePattern]"] +HeadersInput = Union[ResourcePattern, "Iterable[ResourcePattern]"] +MethodsInput = Union[str, "Iterable[str]"] +ExposeHeadersInput = Union[str, "Iterable[str]", None] +MaxAgeInput = Union[timedelta, int, str, None] + + +@dataclass(frozen=True) +class _ComputedCorsOptions: + """The fully-resolved CORS options for a single route. + + Private internal representation, built only by :func:`serialize_options`. + Configure CORS via :class:`~flask_cors.CORS` / :func:`cross_origin` keyword + arguments, not by constructing this directly. + """ + + # ``origins``/``allow_headers`` are normalized to a list of patterns: each + # entry is either a compiled ``re.Pattern`` (regex) or a literal ``str``, + # so request-time matching never has to guess which it is. + origins: list[ResourcePattern] + allow_headers: list[ResourcePattern] + # Whether ``origins`` includes the catch-all ``.*`` wildcard. Computed once + # here because the wildcard is otherwise indistinguishable from a regex + # after the origins have been compiled. + allow_all_origins: bool + # ``methods``/``expose_headers`` join to a comma-separated string, or are + # ``None`` when not configured. + methods: str | None + expose_headers: str | None + # ``max_age`` is left as-is unless a timedelta is given, then stringified. + max_age: str | int | None + supports_credentials: bool + send_wildcard: bool + automatic_options: bool + vary_header: bool + intercept_exceptions: bool + always_send: bool + allow_private_network: bool + + +class CrossOriginOptionsInput(TypedDict, total=False): + """The keyword options accepted by :func:`cross_origin`. + + Every key is optional (``total=False``): only the options a caller + actually passes are forwarded, so app-level ``CORS_*`` configuration and + the defaults still apply to anything left out. Used with PEP 692 + ``Unpack`` to type ``**kwargs`` precisely. + """ + + origins: OriginsInput + methods: MethodsInput + expose_headers: ExposeHeadersInput + allow_headers: HeadersInput + supports_credentials: bool + max_age: MaxAgeInput + send_wildcard: bool + vary_header: bool + automatic_options: bool + always_send: bool + allow_private_network: bool + + +class CorsOptionsInput(CrossOriginOptionsInput, total=False): + """The keyword options accepted by :class:`CORS` / :meth:`CORS.init_app`. + + Extends :class:`CrossOriginOptionsInput` with the application-level options that + only make sense for the extension. + """ + + resources: ResourceSpec + intercept_exceptions: bool + + # Response Headers ACL_ORIGIN = "Access-Control-Allow-Origin" ACL_METHODS = "Access-Control-Allow-Methods" @@ -43,77 +136,97 @@ CONFIG_OPTIONS = [ # to a view. FLASK_CORS_EVALUATED = "_FLASK_CORS_EVALUATED" -# Strange, but this gets the type of a compiled regex, which is otherwise not -# exposed in a public API. -RegexObject = type(re.compile("")) -DEFAULT_OPTIONS = dict( - origins="*", - methods=ALL_METHODS, - allow_headers="*", - expose_headers=None, - supports_credentials=False, - max_age=None, - send_wildcard=False, - automatic_options=True, - vary_header=True, - resources=r"/*", - intercept_exceptions=True, - always_send=True, - allow_private_network=False, -) +# The type of a compiled regular expression. Exposed publicly as +# ``re.Pattern`` since Python 3.8. +RegexObject = re.Pattern + +# Characters commonly found in regular expressions. Their presence is used as +# a heuristic for whether a string is a regex rather than a literal value. +_REGEX_HINT_CHARS = frozenset("*\\]?$^[]()") + +# The set of recognized option names, derived from the input schema (which +# lists exactly the keys a user may pass; the resolved dataclass additionally +# carries computed fields such as ``allow_all_origins``). +_KNOWN_OPTIONS = frozenset(CorsOptionsInput.__required_keys__) | frozenset(CorsOptionsInput.__optional_keys__) + +# The raw, un-normalized option defaults. These are merged with app config and +# user kwargs before being handed to :func:`serialize_options`, which produces +# the strongly-typed :class:`_ComputedCorsOptions`. Typing this as ``CorsOptionsInput`` +# lets mypy verify the defaults match the documented input schema. +DEFAULT_OPTIONS: CorsOptionsInput = { + "origins": "*", + "methods": ALL_METHODS, + "allow_headers": "*", + "expose_headers": None, + "supports_credentials": False, + "max_age": None, + "send_wildcard": False, + "automatic_options": True, + "vary_header": True, + "resources": r"/*", + "intercept_exceptions": True, + "always_send": True, + "allow_private_network": False, +} -def parse_resources(resources): +def parse_resources(resources: ResourceSpec) -> list[tuple[ResourcePattern, dict[str, Any]]]: if isinstance(resources, dict): # To make the API more consistent with the decorator, allow a # resource of '*', which is not actually a valid regexp. - resources = [(re_fix(k), v) for k, v in resources.items()] + resource_pairs = [(re_fix(k), v) for k, v in resources.items()] # Sort patterns with static (literal) paths first, then by regex specificity - def sort_key(pair): + def sort_key(pair: tuple[ResourcePattern, Any]) -> tuple[int, int, int, int]: pattern, _ = pair - if isinstance(pattern, RegexObject): + if isinstance(pattern, re.Pattern): return (1, 0, -pattern.pattern.count("/"), -len(pattern.pattern)) elif probably_regex(pattern): return (1, 1, -pattern.count("/"), -len(pattern)) else: return (0, 0, -pattern.count("/"), -len(pattern)) - return sorted(resources, key=sort_key) + resource_pairs.sort(key=sort_key) elif isinstance(resources, str): - return [(re_fix(resources), {})] + resource_pairs = [(re_fix(resources), {})] elif isinstance(resources, Iterable): - return [(re_fix(r), {}) for r in resources] + resource_pairs = [(re_fix(r), {}) for r in resources] # Type of compiled regex is not part of the public API. Test for this # at runtime. - elif isinstance(resources, RegexObject): - return [(re_fix(resources), {})] + elif isinstance(resources, re.Pattern): + resource_pairs = [(re_fix(resources), {})] else: - raise ValueError("Unexpected value for resources argument.") + # ValueError (rather than the arguably more correct TypeError) is the + # historical public contract here; callers may catch it. + raise ValueError("Unexpected value for resources argument.") # noqa: TRY004 + # Resolve each path pattern once (paths are matched case-sensitively), so + # request handling never re-runs the regex heuristic. Patterns are sorted + # above, before compilation, so specificity ordering is preserved. + return [(_resolve_pattern(pattern, ignore_case=False), opts) for (pattern, opts) in resource_pairs] -def get_regexp_pattern(regexp): + +def get_regexp_pattern(regexp: ResourcePattern) -> str: """ Helper that returns regexp pattern from given value. :param regexp: regular expression to stringify - :type regexp: _sre.SRE_Pattern or str + :type regexp: re.Pattern or str :returns: string representation of given regexp pattern :rtype: str """ - try: + if isinstance(regexp, re.Pattern): return regexp.pattern - except AttributeError: - return str(regexp) + return str(regexp) -def get_cors_origins(options, request_origin): - origins = options.get("origins") - wildcard = r".*" in origins +def get_cors_origins(options: _ComputedCorsOptions, request_origin: str | None) -> list[str] | None: + origins = options.origins + wildcard = options.allow_all_origins # If the Origin header is not present terminate this set of steps. # The request is outside the scope of this specification.-- W3Spec @@ -121,7 +234,7 @@ def get_cors_origins(options, request_origin): LOG.debug("CORS request received with 'Origin' %s", request_origin) # If the allowed origins is an asterisk or 'wildcard', always match - if wildcard and options.get("send_wildcard"): + if wildcard and options.send_wildcard: LOG.debug("Allowed origins are set to '*'. Sending wildcard CORS header.") return ["*"] # If the value of the Origin header is a case-insensitive match @@ -139,7 +252,7 @@ def get_cors_origins(options, request_origin): LOG.debug("The request's Origin header does not match any of allowed origins.") return None - elif options.get("always_send"): + elif options.always_send: if wildcard: # If wildcard is in the origins, even if 'send_wildcard' is False, # simply send the wildcard. Unless supports_credentials is True, @@ -147,13 +260,14 @@ def get_cors_origins(options, request_origin): # It is the most-likely to be correct thing to do (the only other # option is to return nothing, which almost certainly not what # the developer wants if the '*' origin was specified. - if options.get("supports_credentials"): + if options.supports_credentials: return None else: return ["*"] else: - # Return all origins that are not regexes. - return sorted([o for o in origins if not probably_regex(o)]) + # Return all literal (non-regex) origins. After resolution a plain + # ``str`` entry is always a literal; regexes are compiled patterns. + return sorted(o for o in origins if isinstance(o, str)) # Terminate these steps, return the original request untouched. else: @@ -163,38 +277,46 @@ def get_cors_origins(options, request_origin): return None -def get_allow_headers(options, acl_request_headers): +def get_allow_headers(options: _ComputedCorsOptions, acl_request_headers: str | None) -> str | None: if acl_request_headers: + allow_headers = options.allow_headers request_headers = [h.strip() for h in acl_request_headers.split(",")] # any header that matches in the allow_headers - matching_headers = filter(lambda h: try_match_any_pattern(h, options.get("allow_headers"), caseSensitive=False), request_headers) + matching_headers = filter( + lambda h: try_match_any_pattern(h, allow_headers, caseSensitive=False), + request_headers, + ) return ", ".join(sorted(matching_headers)) return None -def get_cors_headers(options, request_headers, request_method): +def get_cors_headers( + options: _ComputedCorsOptions, request_headers: Headers, request_method: str +) -> MultiDict[str, str | int | None]: origins_to_set = get_cors_origins(options, request_headers.get("Origin")) - headers = MultiDict() + # Values are header values (str), with ``max_age`` allowing int and a few + # options allowing None; the trailing comprehension drops the None/empty + # entries before the result is returned. + headers: MultiDict[str, str | int | None] = MultiDict() if not origins_to_set: # CORS is not enabled for this route return headers - for origin in origins_to_set: - headers.add(ACL_ORIGIN, origin) + headers.setlist(ACL_ORIGIN, origins_to_set) - headers[ACL_EXPOSE_HEADERS] = options.get("expose_headers") + headers[ACL_EXPOSE_HEADERS] = options.expose_headers - if options.get("supports_credentials"): + if options.supports_credentials: headers[ACL_CREDENTIALS] = "true" # case sensitive if ( ACL_REQUEST_HEADER_PRIVATE_NETWORK in request_headers and request_headers.get(ACL_REQUEST_HEADER_PRIVATE_NETWORK) == "true" ): - allow_private_network = "true" if options.get("allow_private_network") else "false" + allow_private_network = "true" if options.allow_private_network else "false" headers[ACL_RESPONSE_PRIVATE_NETWORK] = allow_private_network # This is a preflight request @@ -204,36 +326,35 @@ def get_cors_headers(options, request_headers, request_method): # If there is no Access-Control-Request-Method header or if parsing # failed, do not set any additional headers - if acl_request_method and acl_request_method in options.get("methods"): + methods = options.methods + if acl_request_method and methods and acl_request_method in methods: # If method is not a case-sensitive match for any of the values in # list of methods do not set any additional headers and terminate # this set of steps. headers[ACL_ALLOW_HEADERS] = get_allow_headers(options, request_headers.get(ACL_REQUEST_HEADERS)) - headers[ACL_MAX_AGE] = options.get("max_age") - headers[ACL_METHODS] = options.get("methods") + headers[ACL_MAX_AGE] = options.max_age + headers[ACL_METHODS] = methods else: LOG.info( "The request's Access-Control-Request-Method header does not match allowed methods. CORS headers will not be applied." ) # http://www.w3.org/TR/cors/#resource-implementation - if options.get("vary_header"): - # Only set header if the origin returned will vary dynamically, - # i.e. if we are not returning an asterisk, and there are multiple - # origins that can be matched. - if headers[ACL_ORIGIN] == "*": - pass - elif ( - len(options.get("origins")) > 1 - or len(origins_to_set) > 1 - or any(map(probably_regex, options.get("origins"))) + if options.vary_header: + # Only set the header if the resolved origin can vary dynamically: + # i.e. we are not returning a wildcard, and more than one origin (or a + # regex) could have matched. + origins = options.origins + returns_wildcard = headers[ACL_ORIGIN] == "*" + if not returns_wildcard and ( + len(origins) > 1 or len(origins_to_set) > 1 or any(isinstance(o, re.Pattern) for o in origins) ): headers.add("Vary", "Origin") return MultiDict((k, v) for k, v in headers.items() if v) -def set_cors_headers(resp, options): +def set_cors_headers(resp: Response, options: _ComputedCorsOptions) -> Response: """ Performs the actual evaluation of Flask-CORS options and actually modifies the response object. @@ -251,6 +372,8 @@ def set_cors_headers(resp, options): # objects (Werkzeug Headers work as well). This is a problem because # headers allow repeated values. if not isinstance(resp.headers, Headers) and not isinstance(resp.headers, MultiDict): + # The non-standard header container is replaced with a MultiDict so + # repeated headers behave consistently. resp.headers = MultiDict(resp.headers) headers_to_set = get_cors_headers(options, request.headers, request.method) @@ -263,17 +386,16 @@ def set_cors_headers(resp, options): return resp -def probably_regex(maybe_regex): - if isinstance(maybe_regex, RegexObject): +def probably_regex(maybe_regex: ResourcePattern) -> bool: + if isinstance(maybe_regex, re.Pattern): return True - else: - common_regex_chars = ["*", "\\", "]", "?", "$", "^", "[", "]", "(", ")"] - # Use common characters used in regular expressions as a proxy - # for if this string is in fact a regex. - return any(c in maybe_regex for c in common_regex_chars) + + # Use characters common in regular expressions as a proxy for whether + # this string is in fact a regex. + return not _REGEX_HINT_CHARS.isdisjoint(maybe_regex) -def re_fix(reg): +def re_fix(reg: ResourcePattern) -> ResourcePattern: """ Replace the invalid regex r'*' with the valid, wildcard regex r'/.*' to enable the CORS app extension to have a more user friendly api. @@ -281,58 +403,57 @@ def re_fix(reg): return r".*" if reg == r"*" else reg -def try_match_any_pattern(inst, patterns, caseSensitive=True): +def try_match_any_pattern(inst: str, patterns: Iterable[ResourcePattern], caseSensitive: bool = True) -> bool: return any(try_match_pattern(inst, pattern, caseSensitive) for pattern in patterns) -def try_match_pattern(value, pattern, caseSensitive=True): + +def try_match_pattern(value: Any, pattern: ResourcePattern, caseSensitive: bool = True) -> bool: """ - Safely attempts to match a pattern or string to a value. This - function can be used to match request origins, headers, or paths. - The value of caseSensitive should be set in accordance to the - data being compared e.g. origins and headers are case insensitive - whereas paths are case-sensitive + Match a value against a *resolved* pattern: a compiled ``re.Pattern`` + (whose case-sensitivity is already baked in) or a literal string. Used to + match request origins, headers, or paths; ``caseSensitive`` only affects + the literal comparison (origins and headers are case insensitive, paths are + case sensitive). """ - if isinstance(pattern, RegexObject): - return re.match(pattern, value) - if probably_regex(pattern): - flags = 0 if caseSensitive else re.IGNORECASE - try: - return re.match(pattern, value, flags=flags) - except re.error: - return False + if isinstance(pattern, re.Pattern): + return bool(pattern.match(value)) try: v = str(value) p = str(pattern) - return v == p if caseSensitive else v.casefold() == p.casefold() except Exception: - return value == pattern + return bool(value == pattern) + return v == p if caseSensitive else v.casefold() == p.casefold() + -def get_cors_options(appInstance, *dicts): +def merge_options(appInstance: Flask | Blueprint | None, *dicts: Mapping[str, Any]) -> dict[str, Any]: """ - Compute CORS options for an application by combining the DEFAULT_OPTIONS, - the app's configuration-specified options and any dictionaries passed. The - last specified option wins. + Merge the DEFAULT_OPTIONS, the app's configuration-specified options and any + dictionaries passed into a single raw options mapping. The last specified + option wins. """ - options = DEFAULT_OPTIONS.copy() + options: dict[str, Any] = dict(DEFAULT_OPTIONS) options.update(get_app_kwarg_dict(appInstance)) - if dicts: - for d in dicts: - options.update(d) + for d in dicts: + options.update(d) + return options - return serialize_options(options) +def get_cors_options(appInstance: Flask | Blueprint | None, *dicts: Mapping[str, Any]) -> _ComputedCorsOptions: + """Compute the resolved CORS options for an application (see merge_options).""" + return serialize_options(merge_options(appInstance, *dicts)) -def get_app_kwarg_dict(appInstance=None): + +def get_app_kwarg_dict(appInstance: Flask | Blueprint | None = None) -> dict[str, Any]: """Returns the dictionary of CORS specific app configurations.""" app = appInstance or current_app # In order to support blueprints which do not have a config attribute app_config = getattr(app, "config", {}) - return {k.lower().replace("cors_", ""): app_config.get(k) for k in CONFIG_OPTIONS if app_config.get(k) is not None} + return {k.lower().replace("cors_", ""): value for k in CONFIG_OPTIONS if (value := app_config.get(k)) is not None} -def flexible_str(obj): +def flexible_str(obj: Any) -> str | None: """ A more flexible str function which intelligently handles stringifying strings, lists and other iterables. The results are lexographically sorted @@ -347,53 +468,98 @@ def flexible_str(obj): return str(obj) -def serialize_option(options_dict, key, upper=False): - if key in options_dict: - value = flexible_str(options_dict[key]) - options_dict[key] = value.upper() if upper else value - - -def ensure_iterable(inst): +def ensure_iterable(inst: Any) -> Iterable[Any]: """ Wraps scalars or string types as a list, or returns the iterable instance. """ if isinstance(inst, str) or not isinstance(inst, Iterable): return [inst] else: - return inst + return cast("Iterable[Any]", inst) -def sanitize_regex_param(param): +def sanitize_regex_param(param: Any) -> list[ResourcePattern]: return [re_fix(x) for x in ensure_iterable(param)] -def serialize_options(opts): +def _resolve_pattern(pattern: ResourcePattern, *, ignore_case: bool) -> ResourcePattern: + """Resolve a single raw pattern once, at configuration time. + + A regex-like string is compiled to a ``re.Pattern`` (with the correct + case-sensitivity baked in); a literal string and an already-compiled + pattern are returned unchanged. An invalid regular expression raises a + ``ValueError`` so the misconfiguration fails fast at setup time rather than + silently never matching. """ - A helper method to serialize and processes the options dictionary. + if isinstance(pattern, str) and probably_regex(pattern): + try: + return re.compile(pattern, re.IGNORECASE if ignore_case else 0) + except re.error as exc: + raise ValueError(f"Invalid regular expression in Flask-CORS configuration: {pattern!r}") from exc + return pattern + + +def _resolve_patterns(patterns: Iterable[ResourcePattern], *, ignore_case: bool) -> list[ResourcePattern]: + """Resolve a list of patterns (see :func:`_resolve_pattern`). + + After this, the request path can simply test ``isinstance(p, re.Pattern)`` + instead of re-running the :func:`probably_regex` heuristic on every request. """ - options = (opts or {}).copy() + return [_resolve_pattern(p, ignore_case=ignore_case) for p in patterns] + - for key in opts.keys(): - if key not in DEFAULT_OPTIONS: +def serialize_options(opts: Mapping[str, Any]) -> _ComputedCorsOptions: + """ + Normalize a raw options mapping into a strongly-typed :class:`_ComputedCorsOptions`. + + This is the single boundary where loosely-typed user input (arbitrary + keyword arguments, app config, resource dictionaries) is parsed into the + concrete, per-field types the rest of the package relies on. + """ + for key in opts: + if key not in _KNOWN_OPTIONS: LOG.warning("Unknown option passed to Flask-CORS: %s", key) - # Ensure origins is a list of allowed origins with at least one entry. - options["origins"] = sanitize_regex_param(options.get("origins")) - options["allow_headers"] = sanitize_regex_param(options.get("allow_headers")) + # Ensure origins is a list of allowed origins with at least one entry. The + # wildcard is detected before the patterns are compiled (afterwards ``.*`` + # is just another compiled regex, indistinguishable from a real one). + sanitized_origins = sanitize_regex_param(opts.get("origins")) + allow_all_origins = r".*" in sanitized_origins + supports_credentials = bool(opts.get("supports_credentials")) + send_wildcard = bool(opts.get("send_wildcard")) # This is expressly forbidden by the spec. Raise a value error so people # don't get burned in production. - if r".*" in options["origins"] and options["supports_credentials"] and options["send_wildcard"]: + if allow_all_origins and supports_credentials and send_wildcard: raise ValueError( "Cannot use supports_credentials in conjunction with" "an origin string of '*'. See: " "http://www.w3.org/TR/cors/#resource-requests" ) - serialize_option(options, "expose_headers") - serialize_option(options, "methods", upper=True) + origins = _resolve_patterns(sanitized_origins, ignore_case=True) + allow_headers = _resolve_patterns(sanitize_regex_param(opts.get("allow_headers")), ignore_case=True) - if isinstance(options.get("max_age"), timedelta): - options["max_age"] = str(int(options["max_age"].total_seconds())) + methods = flexible_str(opts.get("methods")) + if methods is not None: + methods = methods.upper() - return options + max_age = opts.get("max_age") + if isinstance(max_age, timedelta): + max_age = str(int(max_age.total_seconds())) + + return _ComputedCorsOptions( + origins=origins, + allow_headers=allow_headers, + allow_all_origins=allow_all_origins, + methods=methods, + expose_headers=flexible_str(opts.get("expose_headers")), + max_age=max_age, + supports_credentials=supports_credentials, + send_wildcard=send_wildcard, + automatic_options=bool(opts.get("automatic_options")), + vary_header=bool(opts.get("vary_header")), + intercept_exceptions=bool(opts.get("intercept_exceptions")), + always_send=bool(opts.get("always_send")), + allow_private_network=bool(opts.get("allow_private_network")), + ) diff --git a/contrib/python/Flask-Cors/py3/flask_cors/decorator.py b/contrib/python/Flask-Cors/py3/flask_cors/decorator.py index f8915b62d59..d3c883d4bae 100644 --- a/contrib/python/Flask-Cors/py3/flask_cors/decorator.py +++ b/contrib/python/Flask-Cors/py3/flask_cors/decorator.py @@ -1,14 +1,23 @@ +from __future__ import annotations + import logging +from collections.abc import Callable from functools import update_wrapper +from typing import TYPE_CHECKING, Any, TypeVar, cast + +from flask import Response, current_app, make_response, request -from flask import current_app, make_response, request +from .core import FLASK_CORS_EVALUATED, CrossOriginOptionsInput, get_cors_options, set_cors_headers -from .core import FLASK_CORS_EVALUATED, get_cors_options, set_cors_headers +if TYPE_CHECKING: + from typing_extensions import Unpack LOG = logging.getLogger(__name__) +F = TypeVar("F", bound=Callable[..., Any]) + -def cross_origin(*args, **kwargs): +def cross_origin(*args: Any, **kwargs: Unpack[CrossOriginOptionsInput]) -> Callable[[F], F]: """ This function is the decorator which is used to wrap a Flask route with. In the simplest case, simply use the default parameters to allow all @@ -96,7 +105,7 @@ def cross_origin(*args, **kwargs): """ _options = kwargs - def decorator(f): + def decorator(f: F) -> F: LOG.debug("Enabling %s for cross_origin using options:%s", f, _options) # If True, intercept OPTIONS requests by modifying the view function, @@ -107,15 +116,20 @@ def cross_origin(*args, **kwargs): # decorator (which is actually wraps the function object we return) # intercepts OPTIONS handling, and requests will not have CORS headers if _options.get("automatic_options", True): - f.required_methods = getattr(f, "required_methods", set()) - f.required_methods.add("OPTIONS") - f.provide_automatic_options = False + # ``f`` is a plain view function; these attributes are read by + # Flask's routing. Use an ``Any`` alias so the dynamic attribute + # assignment type-checks. + view_func: Any = f + required_methods = set(getattr(f, "required_methods", set())) + required_methods.add("OPTIONS") + view_func.required_methods = required_methods + view_func.provide_automatic_options = False - def wrapped_function(*args, **kwargs): + def wrapped_function(*args: Any, **kwargs: Any) -> Response: # Handle setting of Flask-Cors parameters options = get_cors_options(current_app, _options) - if options.get("automatic_options") and request.method == "OPTIONS": + if options.automatic_options and request.method == "OPTIONS": resp = current_app.make_default_options_response() else: resp = make_response(f(*args, **kwargs)) @@ -124,6 +138,6 @@ def cross_origin(*args, **kwargs): setattr(resp, FLASK_CORS_EVALUATED, True) return resp - return update_wrapper(wrapped_function, f) + return cast(F, update_wrapper(wrapped_function, f)) return decorator diff --git a/contrib/python/Flask-Cors/py3/flask_cors/extension.py b/contrib/python/Flask-Cors/py3/flask_cors/extension.py index 434f65eaa20..1e172142c43 100644 --- a/contrib/python/Flask-Cors/py3/flask_cors/extension.py +++ b/contrib/python/Flask-Cors/py3/flask_cors/extension.py @@ -1,9 +1,29 @@ +from __future__ import annotations + import logging +from typing import TYPE_CHECKING, Any from urllib.parse import unquote -from flask import request +from flask import Blueprint, Flask, Response, request + +from .core import ( + ACL_ORIGIN, + CorsOptionsInput, + ResourcePattern, + _ComputedCorsOptions, + get_cors_options, + get_regexp_pattern, + merge_options, + parse_resources, + serialize_options, + set_cors_headers, + try_match_pattern, +) + +if TYPE_CHECKING: + from collections.abc import Callable -from .core import ACL_ORIGIN, get_cors_options, get_regexp_pattern, parse_resources, set_cors_headers, try_match_pattern + from typing_extensions import Unpack LOG = logging.getLogger(__name__) @@ -140,50 +160,61 @@ class CORS: :type allow_private_network: bool """ - def __init__(self, app=None, **kwargs): + def __init__(self, app: Flask | Blueprint | None = None, **kwargs: Unpack[CorsOptionsInput]) -> None: self._options = kwargs if app is not None: self.init_app(app, **kwargs) - def init_app(self, app, **kwargs): + def init_app(self, app: Flask | Blueprint, **kwargs: Unpack[CorsOptionsInput]) -> None: # The resources and options may be specified in the App Config, the CORS constructor # or the kwargs to the call to init_app. - options = get_cors_options(app, self._options, kwargs) + merged = merge_options(app, self._options, kwargs) + options = serialize_options(merged) # Flatten our resources into a list of the form - # (pattern_or_regexp, dictionary_of_options) - resources = parse_resources(options.get("resources")) + # (pattern_or_regexp, dictionary_of_raw_options). ``resources`` is a raw + # spec consumed only here, so it is read from the merged mapping rather + # than carried on the resolved options. + resources = parse_resources(merged.get("resources", r"/*")) - # Compute the options for each resource by combining the options from - # the app's configuration, the constructor, the kwargs to init_app, and - # finally the options specified in the resources dictionary. - resources = [(pattern, get_cors_options(app, options, opts)) for (pattern, opts) in resources] + # Compute the options for each resource from the raw inputs (defaults, + # app config, constructor kwargs, init_app kwargs) plus the per-resource + # overrides. Re-merging the raw inputs rather than the already-resolved + # ``options`` avoids re-normalizing compiled regexes. + resolved_resources: list[tuple[ResourcePattern, _ComputedCorsOptions]] = [ + (pattern, get_cors_options(app, self._options, kwargs, opts)) for (pattern, opts) in resources + ] # Create a human-readable form of these resources by converting the compiled # regular expressions into strings. - resources_human = {get_regexp_pattern(pattern): opts for (pattern, opts) in resources} + resources_human = {get_regexp_pattern(pattern): opts for (pattern, opts) in resolved_resources} LOG.debug("Configuring CORS with resources: %s", resources_human) - cors_after_request = make_after_request_function(resources) + cors_after_request = make_after_request_function(resolved_resources) app.after_request(cors_after_request) - # Wrap exception handlers with cross_origin - # These error handlers will still respect the behavior of the route - if options.get("intercept_exceptions", True): + # Wrap exception handlers with cross_origin so error responses also + # receive CORS headers. Blueprints have no ``handle_exception`` / + # ``make_response``, so the ``hasattr`` guard skips them; the ``Any`` + # alias lets us read and reassign those Flask-only methods without + # tripping mypy's attr-defined / method-assign checks for ``Blueprint``. + if options.intercept_exceptions and hasattr(app, "handle_exception"): + app_any: Any = app - def _after_request_decorator(f): - def wrapped_function(*args, **kwargs): - return cors_after_request(app.make_response(f(*args, **kwargs))) + def _after_request_decorator(f: Callable[..., Any]) -> Callable[..., Response]: + def wrapped_function(*args: Any, **kwargs: Any) -> Response: + return cors_after_request(app_any.make_response(f(*args, **kwargs))) return wrapped_function - if hasattr(app, "handle_exception"): - app.handle_exception = _after_request_decorator(app.handle_exception) - app.handle_user_exception = _after_request_decorator(app.handle_user_exception) + app_any.handle_exception = _after_request_decorator(app_any.handle_exception) + app_any.handle_user_exception = _after_request_decorator(app_any.handle_user_exception) -def make_after_request_function(resources): - def cors_after_request(resp): +def make_after_request_function( + resources: list[tuple[ResourcePattern, _ComputedCorsOptions]], +) -> Callable[[Response], Response]: + def cors_after_request(resp: Response) -> Response: # If CORS headers are set in a view decorator, pass if resp.headers is not None and resp.headers.get(ACL_ORIGIN): LOG.debug("CORS have been already evaluated, skipping") diff --git a/contrib/python/Flask-Cors/py3/flask_cors/py.typed b/contrib/python/Flask-Cors/py3/flask_cors/py.typed new file mode 100644 index 00000000000..e69de29bb2d --- /dev/null +++ b/contrib/python/Flask-Cors/py3/flask_cors/py.typed diff --git a/contrib/python/Flask-Cors/py3/ya.make b/contrib/python/Flask-Cors/py3/ya.make index 4c888cf5fde..4e0f3d25720 100644 --- a/contrib/python/Flask-Cors/py3/ya.make +++ b/contrib/python/Flask-Cors/py3/ya.make @@ -2,7 +2,7 @@ PY3_LIBRARY() -VERSION(6.0.3) +VERSION(6.0.5) LICENSE(MIT) @@ -25,6 +25,7 @@ RESOURCE_FILES( PREFIX contrib/python/Flask-Cors/py3/ .dist-info/METADATA .dist-info/top_level.txt + flask_cors/py.typed ) END() diff --git a/contrib/python/clickhouse-connect/.dist-info/METADATA b/contrib/python/clickhouse-connect/.dist-info/METADATA index cb9928f01a9..cf511761a8d 100644 --- a/contrib/python/clickhouse-connect/.dist-info/METADATA +++ b/contrib/python/clickhouse-connect/.dist-info/METADATA @@ -1,6 +1,6 @@ Metadata-Version: 2.4 Name: clickhouse-connect -Version: 1.1.1 +Version: 1.3.0 Summary: ClickHouse Database Core Driver for Python, Pandas, and Superset Home-page: https://github.com/ClickHouse/clickhouse-connect Author: ClickHouse Inc. diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/_version.py b/contrib/python/clickhouse-connect/clickhouse_connect/_version.py index 64c8c21b112..d28b3ddc3b1 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/_version.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/_version.py @@ -1 +1 @@ -version = "1.1.1" +version = "1.3.0" diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/alembic/impl.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/alembic/impl.py index 771692ad015..1f942979c07 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/alembic/impl.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/alembic/impl.py @@ -55,11 +55,20 @@ class ClickHouseImpl(DefaultImpl): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + self._add_integration_tag() if self.context_opts.get("include_schemas") and not self.context_opts.get("version_table_schema") and self.connection is not None: current_database = self.connection.execute(text("SELECT currentDatabase()")).scalar() if current_database: self.context_opts["version_table_schema"] = current_database + def _add_integration_tag(self) -> None: + if self.connection is None: + return + try: + self.connection.connection.driver_connection.client._add_integration_tag("alembic") + except Exception: + pass + def version_table_impl( self, *, diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/datatypes/sqltypes.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/datatypes/sqltypes.py index 94d8e67fb40..5356150545b 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/datatypes/sqltypes.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/datatypes/sqltypes.py @@ -1,8 +1,18 @@ +import ipaddress +import uuid from collections.abc import Sequence from enum import Enum as PyEnum from sqlalchemy.exc import ArgumentError from sqlalchemy.types import ( + ARRAY, + Float, + Integer, + Interval, + Numeric, + UserDefinedType, +) +from sqlalchemy.types import ( Boolean as SqlaBoolean, ) from sqlalchemy.types import ( @@ -12,17 +22,10 @@ from sqlalchemy.types import ( DateTime as SqlaDateTime, ) from sqlalchemy.types import ( - Float, - Integer, - Interval, - Numeric, - UserDefinedType, -) -from sqlalchemy.types import ( String as SqlaString, ) -from clickhouse_connect.cc_sqlalchemy.datatypes.base import ChSqlaType, schema_types +from clickhouse_connect.cc_sqlalchemy.datatypes.base import ChSqlaType, schema_types, sqla_type_from_name from clickhouse_connect.datatypes.base import EMPTY_TYPE_DEF, LC_TYPE_DEF, NULLABLE_TYPE_DEF, TypeDef from clickhouse_connect.datatypes.numeric import Enum8 as ChEnum8 from clickhouse_connect.datatypes.numeric import Enum16 as ChEnum16 @@ -209,43 +212,43 @@ class FixedString(ChSqlaType, SqlaString): class IPv4(ChSqlaType, UserDefinedType): - python_type = None + python_type = ipaddress.IPv4Address class IPv6(ChSqlaType, UserDefinedType): - python_type = None + python_type = ipaddress.IPv6Address class UUID(ChSqlaType, UserDefinedType): - python_type = None + python_type = uuid.UUID class Nothing(ChSqlaType, UserDefinedType): - python_type = None + python_type = type(None) class Point(ChSqlaType, UserDefinedType): - python_type = None + python_type = tuple class Ring(ChSqlaType, UserDefinedType): - python_type = None + python_type = list class Polygon(ChSqlaType, UserDefinedType): - python_type = None + python_type = list class MultiPolygon(ChSqlaType, UserDefinedType): - python_type = None + python_type = list class LineString(ChSqlaType, UserDefinedType): - python_type = None + python_type = list class MultiLineString(ChSqlaType, UserDefinedType): - python_type = None + python_type = list class Date(ChSqlaType, SqlaDate): @@ -412,8 +415,9 @@ class LowCardinality: return element.__class__(type_def=TypeDef(wrappers, orig.keys, orig.values)) -class Array(ChSqlaType, UserDefinedType): +class Array(ChSqlaType, ARRAY): python_type = list + dimensions = 1 def __init__(self, element: ChSqlaType | type[ChSqlaType] = None, type_def: TypeDef = None): """ @@ -425,7 +429,12 @@ class Array(ChSqlaType, UserDefinedType): if callable(element): element = element() type_def = TypeDef(values=(element.name,)) - super().__init__(type_def) + ChSqlaType.__init__(self, type_def) + # Set item_type directly; calling ARRAY.__init__ would reject nested Array(Array(T)), + # which CH supports natively (CH expresses dimensions via nesting, not a dim count). + # as_tuple has no class-level default, so set it here to satisfy ARRAY result processing. + self.item_type = sqla_type_from_name(type_def.values[0]) + self.as_tuple = False class Map(ChSqlaType, UserDefinedType): @@ -489,7 +498,7 @@ class JSON(ChSqlaType, UserDefinedType): Note this isn't currently supported for insert/select, only table definitions """ - python_type = None + python_type = dict class Nested(ChSqlaType, UserDefinedType): @@ -497,11 +506,11 @@ class Nested(ChSqlaType, UserDefinedType): Note this isn't currently supported for insert/select, only table definitions """ - python_type = None + python_type = list class SimpleAggregateFunction(ChSqlaType, UserDefinedType): - python_type = None + python_type = str def __init__( self, @@ -532,7 +541,7 @@ class AggregateFunction(ChSqlaType, UserDefinedType): Note this isn't currently supported for insert/select, only table definitions """ - python_type = None + python_type = str def __init__(self, *params, type_def: TypeDef = None): """ diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/dialect.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/dialect.py index 46eec9166bd..368e80422ba 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/dialect.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/dialect.py @@ -5,7 +5,7 @@ from sqlalchemy.exc import NoResultFound, NoSuchTableError from clickhouse_connect import dbapi from clickhouse_connect.cc_sqlalchemy import dialect_name, ischema_names -from clickhouse_connect.cc_sqlalchemy.inspector import ChInspector, get_table_metadata +from clickhouse_connect.cc_sqlalchemy.inspector import ChInspector, get_columns, get_table_metadata from clickhouse_connect.cc_sqlalchemy.sql import full_table from clickhouse_connect.cc_sqlalchemy.sql.compiler import ChStatementCompiler from clickhouse_connect.cc_sqlalchemy.sql.ddlcompiler import ChDDLCompiler @@ -61,6 +61,11 @@ class ClickHouseDialect(DefaultDialect): ), ] + def __init__(self, server_side_params: bool = False, **kwargs): + # Set before super().__init__() so ChIdentifierPreparer can read it when built. + self.server_side_params = server_side_params + super().__init__(**kwargs) + # SQA 1 compatibility @classmethod @@ -89,6 +94,9 @@ class ClickHouseDialect(DefaultDialect): cmd += " FROM " + quote_identifier(schema) return [row.name for row in connection.execute(text(cmd))] + def get_columns(self, connection, table_name, schema=None, **kw): + return get_columns(connection, table_name, schema) + def get_primary_keys(self, connection, table_name, schema=None, **kw): return [] diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/inspector.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/inspector.py index 067c9fa029c..e566fec988c 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/inspector.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/inspector.py @@ -4,7 +4,7 @@ from collections.abc import Collection from typing import Any import sqlalchemy.schema as sa_schema -from sqlalchemy import text +from sqlalchemy import String, bindparam, text from sqlalchemy.engine.reflection import Inspector from sqlalchemy.exc import NoResultFound @@ -27,7 +27,9 @@ def _database_name(connection, schema: str | None) -> str: def get_table_metadata(connection, table_name, schema=None): database = _database_name(connection, schema) result_set = connection.execute( - text("SELECT engine, engine_full, comment FROM system.tables WHERE database = :database AND name = :table_name"), + text("SELECT engine, engine_full, comment FROM system.tables WHERE database = :database AND name = :table_name").bindparams( + bindparam("database", type_=String()), bindparam("table_name", type_=String()) + ), {"database": database, "table_name": table_name}, ) row = next(result_set, None) @@ -119,6 +121,36 @@ def get_dictionary_metadata(connection, table_name: str, schema: str | None = No return metadata +def get_columns(connection, table_name: str, schema: str | None = None) -> list[dict[str, Any]]: + table_metadata = get_table_metadata(connection, table_name, schema) + if table_metadata.engine == "Dictionary": + return get_dictionary_columns(connection, table_name, schema) + table_id = full_table(table_name, schema) + result_set = connection.execute(text(f"DESCRIBE TABLE {table_id}")) + if not result_set: + raise NoResultFound(f"Table {table_id} does not exist") + columns = [] + for row in result_set: + sqla_type = sqla_type_from_name(row.type.replace("\n", "")) + col = { + "name": row.name, + "type": sqla_type, + "nullable": sqla_type.nullable, + "autoincrement": False, + "comment": row.comment or None, + "clickhouse_codec": row.codec_expression or None, + "clickhouse_ttl": text(row.ttl_expression) if row.ttl_expression else None, + } + if row.default_type == "DEFAULT" and row.default_expression: + col["server_default"] = text(row.default_expression) + elif row.default_type == "MATERIALIZED" and row.default_expression: + col["clickhouse_materialized"] = text(row.default_expression) + elif row.default_type == "ALIAS" and row.default_expression: + col["clickhouse_alias"] = text(row.default_expression) + columns.append(col) + return columns + + class ChInspector(Inspector): def reflect_table( self, @@ -155,30 +187,4 @@ class ChInspector(Inspector): table.kwargs["clickhouse_engine"] = table.engine def get_columns(self, table_name, schema=None, **_kwargs): - table_metadata = get_table_metadata(self.bind, table_name, schema) - if table_metadata.engine == "Dictionary": - return get_dictionary_columns(self.bind, table_name, schema) - table_id = full_table(table_name, schema) - result_set = self.bind.execute(text(f"DESCRIBE TABLE {table_id}")) - if not result_set: - raise NoResultFound(f"Table {table_id} does not exist") - columns = [] - for row in result_set: - sqla_type = sqla_type_from_name(row.type.replace("\n", "")) - col = { - "name": row.name, - "type": sqla_type, - "nullable": sqla_type.nullable, - "autoincrement": False, - "comment": row.comment or None, - "clickhouse_codec": row.codec_expression or None, - "clickhouse_ttl": text(row.ttl_expression) if row.ttl_expression else None, - } - if row.default_type == "DEFAULT" and row.default_expression: - col["server_default"] = text(row.default_expression) - elif row.default_type == "MATERIALIZED" and row.default_expression: - col["clickhouse_materialized"] = text(row.default_expression) - elif row.default_type == "ALIAS" and row.default_expression: - col["clickhouse_alias"] = text(row.default_expression) - columns.append(col) - return columns + return get_columns(self.bind, table_name, schema) diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/sql/compiler.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/sql/compiler.py index 9dc41b21154..28a86cade19 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/sql/compiler.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/sql/compiler.py @@ -1,10 +1,16 @@ +import re + from sqlalchemy.exc import CompileError from sqlalchemy.sql import elements, sqltypes from sqlalchemy.sql.compiler import SQLCompiler +from sqlalchemy.util import memoized_property from clickhouse_connect.cc_sqlalchemy.datatypes.base import ChSqlaType from clickhouse_connect.cc_sqlalchemy.sql import format_table +# The driver's external_bind_re only recognizes \w+ placeholder names. +_bind_name_re = re.compile(r"\w+\Z") + def _find_outermost_marker(text, markers): """Earliest index in `text` where any of `markers` appears at paren depth 0, skipping @@ -40,12 +46,8 @@ def _find_outermost_marker(text, markers): return -1 -def _resolve_ch_type_name(sqla_type): - """Resolve a SQLAlchemy type instance to a ClickHouse type name string. - - Handles both native ChSqlaType instances which carry their ClickHouse name - directly and generic SQLAlchemy types by mapping to reasonable ClickHouse defaults. - """ +def _ch_type_name(sqla_type): + """Map a SQLAlchemy type to a ClickHouse type name, or None if unmapped.""" if isinstance(sqla_type, ChSqlaType): return sqla_type.name # Order matters so we need to check subtypes before parent types @@ -69,10 +71,93 @@ def _resolve_ch_type_name(sqla_type): return "Date" if isinstance(sqla_type, sqltypes.String): return "String" - return "String" + return None + + +def _resolve_ch_type_name(sqla_type): + """ClickHouse type name for a VALUES column, falling back to String.""" + return _ch_type_name(sqla_type) or "String" + + +def _resolve_ch_bind_type(sqla_type): + """ClickHouse type name for a server-side bind, raising on unmapped types.""" + if isinstance(sqla_type, sqltypes.TupleType): + return f"Tuple({', '.join(_resolve_ch_bind_type(t) for t in sqla_type.types)})" + name = _ch_type_name(sqla_type) + if name is None: + raise CompileError(f"server_side_params needs an explicit type for every bind; none resolved for {sqla_type!r}.") + return name class ChStatementCompiler(SQLCompiler): + # SQLAlchemy 1.4 does not pass bindparam_type to bindparam_string, so stash it here. + _ch_bind_type = None + + @property + def _server_side_params(self): + return getattr(self.dialect, "server_side_params", False) + + @property + def _ch_array_binds(self): + return self.__dict__.setdefault("_ch_array_bind_names", set()) + + @memoized_property + def _bind_processors(self): + """Bind processors with array-bound params dropped; the driver formats those.""" + processors = SQLCompiler._bind_processors.fget(self) + if self._ch_array_binds: + return {k: v for k, v in processors.items() if k not in self._ch_array_binds} + return processors + + def _ch_check_bind_name(self, name): + if not _bind_name_re.match(name): + raise CompileError( + f"server_side_params cannot bind parameter {name!r}: ClickHouse server-side " + "parameter names must match [A-Za-z0-9_]. Rename the bind parameter." + ) + + def visit_bindparam(self, bindparam, **kw): + if not self._server_side_params: + return super().visit_bindparam(bindparam, **kw) + if bindparam.expanding and not kw.get("literal_binds") and not bindparam.literal_execute: + return self._ch_expanding_bindparam(bindparam, **kw) + prev = self._ch_bind_type + self._ch_bind_type = bindparam.type + try: + return super().visit_bindparam(bindparam, **kw) + finally: + self._ch_bind_type = prev + + def _ch_has_bind_processor(self, sqla_type): + if isinstance(sqla_type, sqltypes.TupleType): + return any(self._ch_has_bind_processor(t) for t in sqla_type.types) + return sqla_type._cached_bind_processor(self.dialect) is not None + + def _ch_expanding_bindparam(self, bindparam, **kw): + """Render an IN-family bind as a single {name:Array(Type)} placeholder.""" + # Array binds skip per-element processors, so a type that needs one can't be bound. + if self._ch_has_bind_processor(bindparam.type): + raise CompileError(f"server_side_params cannot bind an IN list of {bindparam.type!r}: it needs a bind processor.") + # Drop the param from post-compile expansion so its list reaches the driver intact. + super().visit_bindparam(bindparam, **kw) + self.post_compile_params = self.post_compile_params.difference([bindparam]) + name = self._truncate_bindparam(bindparam) + actual = self.escaped_bind_names.get(name, name) if self.escaped_bind_names else name + self._ch_check_bind_name(actual) + self._ch_array_binds.add(actual) + return "{" + actual + ":Array(" + _resolve_ch_bind_type(bindparam.type) + ")}" + + def bindparam_string(self, name, **kw): + base = super().bindparam_string(name, **kw) + if not self._server_side_params or kw.get("post_compile") or kw.get("expanding"): + return base + actual = self.escaped_bind_names.get(name, name) if self.escaped_bind_names else name + self._ch_check_bind_name(actual) + ch_type = kw.get("bindparam_type") or self._ch_bind_type + if ch_type is None: + raise CompileError(f"server_side_params requires a typed bind parameter, but none was available for {name!r}.") + return "{" + actual + ":" + _resolve_ch_bind_type(ch_type) + "}" + def _raise_on_escape(self, binary, operator_name: str): if binary.modifiers.get("escape") is not None: raise CompileError(f"ClickHouse does not support the ESCAPE clause on {operator_name}") diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/sql/preparer.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/sql/preparer.py index b22c2b32c61..4521eb97dfa 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/sql/preparer.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/sql/preparer.py @@ -6,5 +6,10 @@ from clickhouse_connect.driver.binding import quote_identifier class ChIdentifierPreparer(IdentifierPreparer): quote_identifier = staticmethod(quote_identifier) + def __init__(self, dialect, **kwargs): + super().__init__(dialect, **kwargs) + if getattr(dialect, "server_side_params", False): + self._double_percents = False + def _requires_quotes(self, _value): return True diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/dbapi/cursor.py b/contrib/python/clickhouse-connect/clickhouse_connect/dbapi/cursor.py index 9ad0431d6ed..af7f36f67b6 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/dbapi/cursor.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/dbapi/cursor.py @@ -1,6 +1,6 @@ import logging import re -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from clickhouse_connect.datatypes.registry import get_from_name from clickhouse_connect.driver import Client @@ -94,10 +94,19 @@ class Cursor: op_columns = None if "VALUES" not in temp.upper(): return False - col_names = list(data[0].keys()) - if op_columns and {unescape_identifier(x) for x in op_columns} != set(col_names): - return False # Data sent in doesn't match the columns in the insert statement - data_values = [list(row.values()) for row in data] + first_row = data[0] + if isinstance(first_row, Mapping): + col_names = list(first_row.keys()) + if op_columns and {unescape_identifier(x) for x in op_columns} != set(col_names): + return False # Data sent in doesn't match the columns in the insert statement + data_values = [list(row.values()) for row in data] + elif isinstance(first_row, Sequence) and not isinstance(first_row, (str, bytes)): + # PEP 249 also allows rows as sequences; take column names from the + # insert statement if present, otherwise insert into all columns + col_names = [unescape_identifier(x) for x in op_columns] if op_columns else "*" + data_values = data + else: + return False self.client.insert(table, data_values, col_names) self.data = [] return True diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/__init__.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/__init__.py index d0fe9f1d8f6..811286ad24c 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/__init__.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/__init__.py @@ -1,8 +1,9 @@ from __future__ import annotations +from collections.abc import Awaitable, Callable from inspect import signature from typing import TYPE_CHECKING, Any -from urllib.parse import parse_qs, urlparse +from urllib.parse import parse_qs, unquote, urlparse import clickhouse_connect.driver.ctypes # noqa: F401 -- side-effect import from clickhouse_connect.driver.client import Client @@ -34,6 +35,11 @@ def default_port(interface: str, secure: bool) -> int: raise ValueError("Unrecognized ClickHouse interface") +def _unquote(value: str | None) -> str | None: + """Percent-decode a DSN component, passing through None/empty.""" + return unquote(value) if value else value + + def _parse_connection_params( host: str | None, username: str | None, @@ -48,12 +54,12 @@ def _parse_connection_params( """Parse and normalize connection parameters including DSN parsing.""" if dsn: parsed = urlparse(dsn) - username = username or parsed.username - password = password or parsed.password + username = username or _unquote(parsed.username) + password = password or _unquote(parsed.password) or "" host = host or parsed.hostname port = port or parsed.port if parsed.path and (not database or database == "__default__"): - database = parsed.path[1:].split("/")[0] + database = unquote(parsed.path[1:].split("/")[0]) database = database or parsed.path for k, v in parse_qs(parsed.query).items(): kwargs[k] = v[0] @@ -75,10 +81,26 @@ def _parse_connection_params( return host, username, password, port, database, interface -def _validate_access_token(access_token: str | None, username: str | None, password: str) -> None: - """Validate that access_token and username/password are not both provided.""" - if access_token and (username or password != ""): - raise ProgrammingError("Cannot use both access_token and username/password") +def _validate_access_token(access_token: str | None, token_provider: Callable[[], str] | None, username: str | None, password: str) -> None: + """Validate that token-based and username/password auth are not mixed.""" + if (access_token or token_provider) and (username or password): + raise ProgrammingError("Cannot use both token authentication and username/password") + if access_token and token_provider: + raise ProgrammingError("Cannot use both access_token and token_provider") + + +def _pop_headers_arg(headers: Any | None, kwargs: dict[str, Any]) -> Any | None: + """Hoist headers parsed through generic kwargs while preserving explicit headers.""" + if "headers" in kwargs: + kwargs_headers = kwargs.pop("headers") + if headers is None: + headers = kwargs_headers + return headers + + +def _validate_headers(headers: Any | None) -> None: + if headers is not None and not isinstance(headers, dict): + raise ProgrammingError("headers must be a dictionary of HTTP header names and values") def create_client( @@ -87,12 +109,14 @@ def create_client( username: str | None = None, password: str = "", access_token: str | None = None, + token_provider: Callable[[], str] | None = None, database: str = "__default__", interface: str | None = None, port: int = 0, secure: bool | str = False, dsn: str | None = None, settings: dict[str, Any] | None = None, + headers: dict[str, str] | None = None, generic_args: dict[str, Any] | None = None, **kwargs, ) -> Client: @@ -106,6 +130,9 @@ def create_client( Should not be set if `access_token` is used. :param access_token: JWT access token (ClickHouse Cloud feature). Should not be set if `username`/`password` are used. + :param token_provider: A callable returning a JWT access token (ClickHouse Cloud feature). Called for the initial token and + again to refresh it whenever the server rejects the current one. + Should not be set if `access_token` or `username`/`password` are used. :param database: The default database for the connection. If not set, ClickHouse Connect will use the default database for username. :param interface: Must be http or https. Defaults to http, or to https if port is set to 8443 or 443 @@ -115,6 +142,9 @@ def create_client( :param dsn: A string in standard DSN (Data Source Name) format. Other connection values (such as host or user) will be extracted from this string if not set otherwise. :param settings: ClickHouse server settings to be used with the session/every request + :param headers: Additional HTTP headers to send with every request. This can be used for proxy or gateway + authentication, such as Cloudflare Access service token headers. These headers are applied after driver defaults, + so they can intentionally override headers such as Authorization or User-Agent. :param generic_args: Used internally to parse DBAPI connection strings into keyword arguments and ClickHouse settings. It is not recommended to use this parameter externally. @@ -154,22 +184,27 @@ def create_client( match the server's column definition which means timezone-aware when the column defines a timezone and naive for bare DateTime columns. :param autogenerate_session_id If set, this will override the 'autogenerate_session_id' common setting. - :param form_encode_query_params If True, query parameters will be sent as form-encoded data in the request body - instead of as URL parameters. This is useful for queries with large parameter sets that might exceed URL length - limits. Only available for query operations (not inserts). Default: False + :param form_encode_query_params If True, always send query parameters as form-encoded data in the request body + instead of as URL parameters. When False, large parameter payloads are still automatically sent as form data to + avoid exceeding URL length limits, except for queries using binary parameter binds, which are only form-encoded + when this is True. Only available for query operations (not inserts). Default: False :return: ClickHouse Connect Client instance """ host, username, password, port, database, interface = _parse_connection_params( host, username, password, port, database, interface, secure, dsn, kwargs ) - _validate_access_token(access_token, username, password) + headers = _pop_headers_arg(headers, kwargs) + _validate_access_token(access_token, token_provider, username, password) settings = settings or {} if interface.startswith("http"): if generic_args: client_params = signature(HttpClient).parameters for name, value in generic_args.items(): - if name in client_params: + if name == "headers": + if headers is None: + headers = value + elif name in client_params: kwargs[name] = value elif name == "compression": if "compress" not in kwargs: @@ -178,7 +213,26 @@ def create_client( if name.startswith("ch_"): name = name[3:] settings[name] = value - return HttpClient(interface, host, port, username, password, database, access_token, settings=settings, **kwargs) + # token auth may also arrive via generic_args (DB-API connect_args); pop both so neither is passed twice + generic_access = kwargs.pop("access_token", None) + generic_token = kwargs.pop("token_provider", None) + access_token = access_token or generic_access + token_provider = token_provider or generic_token + _validate_access_token(access_token, token_provider, username, password) + _validate_headers(headers) + return HttpClient( + interface, + host, + port, + username, + password, + database, + access_token, + token_provider=token_provider, + settings=settings, + headers=headers, + **kwargs, + ) raise ProgrammingError(f"Unrecognized client type {interface}") @@ -188,12 +242,14 @@ async def create_async_client( username: str | None = None, password: str = "", access_token: str | None = None, + token_provider: Callable[[], str | Awaitable[str]] | None = None, database: str = "__default__", interface: str | None = None, port: int = 0, secure: bool | str = False, dsn: str | None = None, settings: dict[str, Any] | None = None, + headers: dict[str, str] | None = None, generic_args: dict[str, Any] | None = None, connector_limit: int = 100, connector_limit_per_host: int = 20, @@ -212,6 +268,9 @@ async def create_async_client( :param username: The ClickHouse username. If not set, the default ClickHouse user will be used. :param password: The password for username. :param access_token: JWT access token. + :param token_provider: A callable returning a JWT access token. Called for the initial token and + again to refresh it whenever the server rejects the current one. Because multiple in-flight requests + may each trigger a refresh concurrently, the callable must be safe to invoke in parallel. :param database: The default database for the connection. If not set, ClickHouse Connect will use the default database for username. :param interface: Must be http or https. Defaults to http, or to https if port is set to 8443 or 443 @@ -221,6 +280,9 @@ async def create_async_client( :param dsn: A string in standard DSN (Data Source Name) format. Other connection values (such as host or user) will be extracted from this string if not set otherwise. :param settings: ClickHouse server settings to be used with the session/every request + :param headers: Additional HTTP headers to send with every request. This can be used for proxy or gateway + authentication, such as Cloudflare Access service token headers. These headers are applied after driver defaults, + so they can intentionally override headers such as Authorization or User-Agent. :param generic_args: Used internally to parse DBAPI connection strings into keyword arguments and ClickHouse settings. It is not recommended to use this parameter externally :param connector_limit: Maximum number of allowable connections to the server @@ -259,9 +321,10 @@ async def create_async_client( match the server's column definition which means timezone-aware when the column defines a timezone and naive for bare DateTime columns. :param autogenerate_session_id If set, this will override the 'autogenerate_session_id' common setting. - :param form_encode_query_params If True, query parameters will be sent as form-encoded data in the request body - instead of as URL parameters. This is useful for queries with large parameter sets that might exceed URL length - limits. Only available for query operations (not inserts). Default: False + :param form_encode_query_params If True, always send query parameters as form-encoded data in the request body + instead of as URL parameters. When False, large parameter payloads are still automatically sent as form data to + avoid exceeding URL length limits, except for queries using binary parameter binds, which are only form-encoded + when this is True. Only available for query operations (not inserts). Default: False :return: ClickHouse Connect AsyncClient instance """ try: @@ -280,13 +343,17 @@ async def create_async_client( host, username, password, port, database, interface = _parse_connection_params( host, username, password, port, database, interface, secure, dsn, kwargs ) - _validate_access_token(access_token, username, password) + headers = _pop_headers_arg(headers, kwargs) + _validate_access_token(access_token, token_provider, username, password) settings = settings or {} if generic_args: client_params = signature(_AsyncClient).parameters for name, value in generic_args.items(): - if name in client_params: + if name == "headers": + if headers is None: + headers = value + elif name in client_params: kwargs[name] = value elif name == "compression": if "compress" not in kwargs: @@ -299,6 +366,13 @@ async def create_async_client( if "autogenerate_session_id" not in kwargs: kwargs["autogenerate_session_id"] = False + # token auth may also arrive via generic_args (DB-API connect_args); pop both so neither is passed twice + generic_access = kwargs.pop("access_token", None) + generic_token = kwargs.pop("token_provider", None) + access_token = access_token or generic_access + token_provider = token_provider or generic_token + _validate_access_token(access_token, token_provider, username, password) + _validate_headers(headers) client = _AsyncClient( interface=interface, host=host, @@ -307,7 +381,9 @@ async def create_async_client( password=password, database=database, access_token=access_token, + token_provider=token_provider, settings=settings, + headers=headers, connector_limit=connector_limit, connector_limit_per_host=connector_limit_per_host, keepalive_timeout=keepalive_timeout, diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/asyncclient.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/asyncclient.py index 1cadab3e907..b0301df84ca 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/asyncclient.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/asyncclient.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import gzip +import inspect import io import json import logging @@ -34,13 +35,20 @@ from clickhouse_connect.datatypes.base import ClickHouseType from clickhouse_connect.datatypes.registry import get_from_name from clickhouse_connect.driver import httputil, options, tzutil from clickhouse_connect.driver.asyncqueue import EOF_SENTINEL, AsyncSyncQueue -from clickhouse_connect.driver.binding import bind_query, quote_identifier +from clickhouse_connect.driver.binding import bind_query, quote_identifier, use_form_encoding from clickhouse_connect.driver.client import Client, _apply_arrow_tz_policy from clickhouse_connect.driver.common import StreamContext, coerce_bool, dict_copy from clickhouse_connect.driver.compression import available_compression from clickhouse_connect.driver.constants import CH_VERSION_WITH_PROTOCOL, PROTOCOL_VERSION_WITH_LOW_CARD from clickhouse_connect.driver.ctypes import RespBuffCls -from clickhouse_connect.driver.exceptions import DatabaseError, DataError, OperationalError, ProgrammingError +from clickhouse_connect.driver.exceptions import ( + DatabaseError, + DataError, + OperationalError, + ProgrammingError, + error_code_from_header, + error_name_from_body, +) from clickhouse_connect.driver.external import ExternalData from clickhouse_connect.driver.insert import InsertContext from clickhouse_connect.driver.models import ColumnDef, SettingDef @@ -60,6 +68,7 @@ from clickhouse_connect.driver.transform import NativeTransform logger = logging.getLogger(__name__) columns_only_re = re.compile(r"LIMIT 0\s*$", re.IGNORECASE) ex_header = "X-ClickHouse-Exception-Code" +auth_failed_ex_code = "516" # ClickHouse AUTHENTICATION_FAILED ex_tag_header = "X-ClickHouse-Exception-Tag" if "br" in available_compression: @@ -197,6 +206,7 @@ class AsyncClient(Client): password: str | None = None, database: str | None = None, access_token: str | None = None, + token_provider: Callable[[], str | Awaitable[str]] | None = None, compress: bool | str = True, connect_timeout: int = 10, send_receive_timeout: int = 300, @@ -224,6 +234,7 @@ class AsyncClient(Client): autogenerate_query_id: bool | None = None, form_encode_query_params: bool = False, rename_response_column: str | None = None, + headers: dict[str, str] | None = None, ): """ Async HTTP Client using aiohttp. Initialization is handled via _initialize(). @@ -243,6 +254,8 @@ class AsyncClient(Client): verify = True tls_mode = tls_mode or "proxy" + self._token_provider = token_provider # initial token is resolved in _initialize() + # Priority: access_token > mutual TLS > basic auth if client_cert and (tls_mode is None or tls_mode == "mutual"): if not username: @@ -329,6 +342,8 @@ class AsyncClient(Client): self._reported_libs = set() self._last_pool_reset = None self.headers["User-Agent"] = self.headers["User-Agent"].replace("mode:sync;", "mode:async;") + if headers: + self.headers.update(headers) # Store aiohttp-specific params for deferred initialization self._compress_param = compress @@ -382,6 +397,9 @@ class AsyncClient(Client): if self._initialized: return + if self._token_provider: + self.set_access_token(await self._resolve_token()) + try: tz_source = self._deferred_tz_source @@ -533,6 +551,15 @@ class AsyncClient(Client): def get_client_setting(self, key) -> str | None: return self._client_settings.get(key) + async def _resolve_token(self) -> str: + # Run sync providers off the event loop; await async providers. + # The provider may be called concurrently if multiple requests get a 516 at the same time; + # it must be safe to invoke in parallel (e.g. if it hits an IdP, consider rate limiting). + result = await asyncio.get_running_loop().run_in_executor(None, self._token_provider) + if inspect.isawaitable(result): + result = await result + return result + def set_access_token(self, access_token: str): auth_header = self.headers.get("Authorization") if auth_header and not auth_header.startswith("Bearer"): @@ -560,13 +587,14 @@ class AsyncClient(Client): context.block_info = True params.update(self._validate_settings(context.settings)) context.rename_response_column = self._rename_response_column + use_form = use_form_encoding(context.final_query, context.bind_params, self.form_encode_query_params) if not context.is_insert and columns_only_re.search(context.uncommented_query): fmt_json_query = f"{context.final_query}\n FORMAT JSON" fields = {"query": fmt_json_query} fields.update(context.bind_params) - if self.form_encode_query_params: + if use_form: files = {} if context.external_data: params.update(context.external_data.query_params) @@ -625,7 +653,7 @@ class AsyncClient(Client): files = None data = None - if self.form_encode_query_params: + if use_form: fields = {"query": final_query} fields.update(context.bind_params) @@ -1105,10 +1133,11 @@ class AsyncClient(Client): files = None body = None - if external_data and not self.form_encode_query_params and isinstance(final_query, bytes): + use_form = use_form_encoding(final_query, bind_params, self.form_encode_query_params) + if external_data and not use_form and isinstance(final_query, bytes): raise ProgrammingError("Binary query cannot be placed in URL when using External Data; enable form encoding.") - if self.form_encode_query_params: + if use_form: files = {} files["query"] = (None, final_query if isinstance(final_query, str) else final_query.decode()) for k, v in bind_params.items(): @@ -1870,26 +1899,29 @@ class AsyncClient(Client): """ try: body = "" + full_body = "" try: raw_body = await response.read() encoding = response.headers.get("Content-Encoding") + loop = asyncio.get_running_loop() if encoding: - loop = asyncio.get_running_loop() def decompress_and_decode(): - decompressed = decompress_response(raw_body, encoding) - return common.format_error(decompressed.decode(errors="backslashreplace")).strip() + return decompress_response(raw_body, encoding).decode(errors="backslashreplace") - body = await loop.run_in_executor(None, decompress_and_decode) + full_body = await loop.run_in_executor(None, decompress_and_decode) else: - loop = asyncio.get_running_loop() - body = await loop.run_in_executor(None, lambda: common.format_error(raw_body.decode(errors="backslashreplace")).strip()) + full_body = await loop.run_in_executor(None, lambda: raw_body.decode(errors="backslashreplace")) + body = common.format_error(full_body).strip() except Exception: logger.warning("Failed to read error response body", exc_info=True) + err_code = response.headers.get(ex_header) + code = error_code_from_header(err_code) + name = error_name_from_body(full_body) if self.show_clickhouse_errors else None + if self.show_clickhouse_errors: - err_code = response.headers.get(ex_header) if err_code: err_str = f"Received ClickHouse exception, code: {err_code}" else: @@ -1905,7 +1937,8 @@ class AsyncClient(Client): finally: response.close() - raise OperationalError(err_str) if retried else DatabaseError(err_str) from None + err_type = OperationalError if retried else DatabaseError + raise err_type(err_str, code=code, name=name) from None async def _raw_request( self, @@ -1950,6 +1983,7 @@ class AsyncClient(Client): req_headers["Host"] = self.server_host_name query_session = final_params.get("session_id") attempts = 0 + auth_retried = False while True: attempts += 1 @@ -2018,6 +2052,17 @@ class AsyncClient(Client): await asyncio.sleep(0.1 * attempts) response.close() continue + if self._token_provider and not auth_retried and response.headers.get(ex_header) == auth_failed_ex_code: + if retry_body is None and not (data is None or isinstance(data, (bytes, bytearray, str, dict))): + await self._error_handler(response) # non-replayable body, surface the auth error instead of retrying + auth_retried = True + self.set_access_token(await self._resolve_token()) + req_headers["Authorization"] = self.headers["Authorization"] + if retry_body is not None: + data = await retry_body() + logger.debug("Refreshing access token after authentication failure") + response.close() + continue await self._error_handler(response) except aiohttp.ClientConnectionError as e: diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/binding.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/binding.py index 3734589a2e2..b4f848dc77a 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/binding.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/binding.py @@ -6,6 +6,7 @@ from collections.abc import Sequence from datetime import date, datetime, timezone, tzinfo from enum import Enum from typing import Any +from urllib.parse import quote, urlencode from clickhouse_connect import common from clickhouse_connect.driver import tzutil @@ -68,22 +69,22 @@ def finalize_query(query: str, parameters: Sequence | dict[str, Any] | None, ser return query % tuple(format_query_value(v, server_tz) for v in parameters) -def _extract_tz_from_type(type_str: str) -> tzinfo | None: - """Extract timezone from a ClickHouse type hint like DateTime64(6, 'UTC'). - - Handles LowCardinality/Nullable wrappers and container types - (Array, Tuple, Map). Returns None if no timezone is found or parsing fails. - """ - try: - base = type_str.strip() - if base.startswith("LowCardinality"): - base = base[15:-1] - if base.startswith("Nullable"): - base = base[9:-1] +def _unwrap_outer(type_str: str) -> tuple[str, tuple]: + """Strip LowCardinality/Nullable wrappers and return (base_name, args)""" + base = type_str.strip() + if base[:15].lower() == "lowcardinality(": + base = base[15:-1] + if base[:9].lower() == "nullable(": + base = base[9:-1] + base_name, values, _ = parse_callable(base) + return base_name, values - base_name, values, _ = parse_callable(base) - if base_name in ("DateTime", "DateTime64"): +def _extract_tz_from_type(type_str: str) -> tzinfo | None: + """Resolve the timezone named in a ClickHouse type hint.""" + try: + base_name, values = _unwrap_outer(type_str) + if base_name.lower() in ("datetime", "datetime64"): for v in values: if isinstance(v, str) and v.startswith("'") and v.endswith("'"): try: @@ -104,6 +105,25 @@ def _extract_tz_from_type(type_str: str) -> tzinfo | None: return None +def _promote_datetime64(type_str: str, value): + """Wrap values bound to a DateTime64 hint in DT64Param to preserve precision.""" + if value is None or "datetime64" not in type_str.lower(): + return value + try: + base_name, values = _unwrap_outer(type_str) + base_name = base_name.lower() + if base_name == "datetime64": + return DT64Param(value) if isinstance(value, datetime) else value + if base_name == "array" and values and isinstance(value, (list, tuple)): + inner = str(values[0]) + return type(value)(_promote_datetime64(inner, x) for x in value) + if base_name == "tuple" and isinstance(value, tuple) and len(values) == len(value): + return tuple(_promote_datetime64(str(t), x) for t, x in zip(values, value)) + return value + except Exception: + return value + + def bind_query( query: str, parameters: Sequence | dict[str, Any] | None, @@ -121,9 +141,13 @@ def bind_query( for key in binary_binds.keys(): del params_copy[key] + matches = external_bind_re.findall(query) + placeholder_names = {name for name, _ in matches} final_params = {} for k, v in params_copy.items(): - if k.endswith("_64"): + # The _64 suffix is a precision hint, not part of the name, unless the + # query binds the full name itself. + if k.endswith("_64") and k not in placeholder_names: if isinstance(v, datetime): k = k[:-3] v = DT64Param(v) @@ -131,7 +155,6 @@ def bind_query( k = k[:-3] v = [DT64Param(x) for x in v] final_params[k] = v - matches = external_bind_re.findall(query) if not matches: query, bound_params = finalize_query(query, final_params, server_tz), {} else: @@ -147,6 +170,7 @@ def bind_query( hint_tz = _extract_tz_from_type(type_str) if hint_tz is not None: tz = hint_tz + v = _promote_datetime64(type_str, v) bound_params[f"param_{k}"] = format_bind_value(v, tz) else: query, bound_params = finalize_query(query, parameters, server_tz), {} @@ -172,6 +196,30 @@ def bind_query( return query, bound_params +# Server-side bind parameters are urlencoded into the request URL. Once the encoded length +# passes this budget the client routes them through multipart form data instead, keeping +# oversized payloads out of the URL where proxies (nginx, ALB, CloudFront) reject them with +# HTTP 414. The threshold leaves ample headroom under common request line limits. +MAX_URL_BIND_PARAM_LENGTH = 4096 + + +def use_form_encoding(query: str | bytes, bind_params: dict[str, str], force_form: bool = False) -> bool: + if force_form: + return True + # Binary binds embed bytes into the query, which the form path cannot round-trip; leave + # those on the default path unless form encoding is explicitly requested. + if isinstance(query, bytes): + return False + if not bind_params: + return False + # Raw length is a lower bound on the encoded length, so large payloads short-circuit + # without materializing the encoded string. + if sum(len(k) + len(str(v)) for k, v in bind_params.items()) > MAX_URL_BIND_PARAM_LENGTH: + return True + # Measure with quote so spaces count as %20, matching the longer of the two client encodings. + return len(urlencode(bind_params, quote_via=quote)) > MAX_URL_BIND_PARAM_LENGTH + + def format_str(value: str): return f"'{escape_str(value)}'" @@ -180,6 +228,10 @@ def escape_str(value: str): return "".join(f"{BS}{c}" if c in must_escape else c for c in value) +def escape_bytes(value): + return "".join(f"{BS}x{b:02x}" for b in value) + + def format_query_value(value: Any, server_tz: tzinfo = timezone.utc): """ Format Python values in a ClickHouse query @@ -191,6 +243,8 @@ def format_query_value(value: Any, server_tz: tzinfo = timezone.utc): return "NULL" if isinstance(value, str): return format_str(value) + if isinstance(value, (bytes, bytearray)): + return f"'{escape_bytes(value)}'" if isinstance(value, DT64Param): return value.format(server_tz, False) if isinstance(value, datetime): @@ -238,6 +292,10 @@ def format_bind_value(value: Any, server_tz: tzinfo = timezone.utc, top_level: b # At the top levels, strings must not be surrounded by quotes return escape_str(value) return format_str(value) + if isinstance(value, (bytes, bytearray)): + if top_level: + return escape_bytes(value) + return f"'{escape_bytes(value)}'" if isinstance(value, DT64Param): return value.format(server_tz, top_level) if isinstance(value, datetime): @@ -261,4 +319,8 @@ def format_bind_value(value: Any, server_tz: tzinfo = timezone.utc, top_level: b return f"{{{', '.join(pairs)}}}" if isinstance(value, Enum): return recurse(value.value) + if isinstance(value, (uuid.UUID, ipaddress.IPv4Address, ipaddress.IPv6Address)): + if top_level: + return str(value) + return f"'{value}'" return str(value) diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/exceptions.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/exceptions.py index 84009a319ab..cddcac55f59 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/exceptions.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/exceptions.py @@ -4,6 +4,26 @@ how useful that naming convention is, but the convention is used for potential i libraries. In most cases docstring are taken from the DBIApi 2.0 documentation """ +import re + +_error_name_re = re.compile(r"\(([A-Z][A-Z0-9_]+)\)") + + +def error_code_from_header(header_value: str | None) -> int | None: + """Parse the numeric ClickHouse error code from the exception header value.""" + if not header_value: + return None + try: + return int(header_value) + except (TypeError, ValueError): + return None + + +def error_name_from_body(body: str | None) -> str | None: + """Extract the symbolic ClickHouse error name (e.g. UNKNOWN_TABLE) from a response body.""" + matches = _error_name_re.findall(body or "") + return matches[-1] if matches else None + class ClickHouseError(Exception): """Exception related to operation with ClickHouse.""" @@ -16,7 +36,16 @@ class Warning(Warning, ClickHouseError): # noqa: N818 class Error(ClickHouseError): """Exception that is the base class of all other error exceptions - (not Warning).""" + (not Warning). + + `code` is the numeric ClickHouse server error code when known, `name` the symbolic + name. Both are None when unavailable, e.g. transport errors or suppressed error detail. + """ + + def __init__(self, *args, code: int | None = None, name: str | None = None): + super().__init__(*args) + self.code = code + self.name = name class InterfaceError(Error): diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/httpclient.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/httpclient.py index 8040cd27c5b..170c2ab6ab8 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/httpclient.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/httpclient.py @@ -19,12 +19,18 @@ from urllib3.response import HTTPResponse from clickhouse_connect import common from clickhouse_connect.datatypes import registry from clickhouse_connect.datatypes.base import ClickHouseType -from clickhouse_connect.driver.binding import bind_query, quote_identifier +from clickhouse_connect.driver.binding import bind_query, quote_identifier, use_form_encoding from clickhouse_connect.driver.client import Client from clickhouse_connect.driver.common import coerce_bool, coerce_int, dict_add, dict_copy from clickhouse_connect.driver.compression import available_compression from clickhouse_connect.driver.ctypes import RespBuffCls -from clickhouse_connect.driver.exceptions import DatabaseError, OperationalError, ProgrammingError +from clickhouse_connect.driver.exceptions import ( + DatabaseError, + OperationalError, + ProgrammingError, + error_code_from_header, + error_name_from_body, +) from clickhouse_connect.driver.external import ExternalData from clickhouse_connect.driver.httputil import ( ResponseSource, @@ -44,6 +50,7 @@ from clickhouse_connect.driver.transform import NativeTransform logger = logging.getLogger(__name__) columns_only_re = re.compile(r"LIMIT 0\s*$", re.IGNORECASE) ex_header = "X-ClickHouse-Exception-Code" +auth_failed_ex_code = "516" # ClickHouse AUTHENTICATION_FAILED ex_tag_header = "X-ClickHouse-Exception-Tag" _REMOTE_CLOSE_ERRORS = (ConnectionResetError, BrokenPipeError) @@ -79,6 +86,7 @@ class HttpClient(Client): password: str, database: str, access_token: str | None = None, + token_provider: Callable[[], str] | None = None, compress: bool | str = True, query_limit: int = 0, query_retries: int = 2, @@ -104,6 +112,7 @@ class HttpClient(Client): proxy_path: str = "", form_encode_query_params: bool = False, rename_response_column: str | None = None, + headers: dict[str, str] | None = None, ): """ Create an HTTP ClickHouse Connect client @@ -150,6 +159,9 @@ class HttpClient(Client): else: self.http = default_pool_manager() + self._token_provider = token_provider + if token_provider: + access_token = token_provider() if access_token: self.headers["Authorization"] = f"Bearer {access_token}" elif (not client_cert or tls_mode in ("strict", "proxy")) and username: @@ -157,6 +169,8 @@ class HttpClient(Client): self._reported_libs = set() self.headers["User-Agent"] = common.build_client_name(client_name) + if headers: + self.headers.update(headers) self._read_format = self._write_format = "Native" self._transform = NativeTransform() @@ -260,10 +274,11 @@ class HttpClient(Client): context.block_info = True params.update(self._validate_settings(context.settings)) context.rename_response_column = self._rename_response_column + use_form = use_form_encoding(context.final_query, context.bind_params, self.form_encode_query_params) if not context.is_insert and columns_only_re.search(context.uncommented_query): # Mirror normal query behavior for form encoding and external data fmt_json_query = f"{context.final_query}\n FORMAT JSON" - if self.form_encode_query_params: + if use_form: fields = {"query": fmt_json_query} fields.update(context.bind_params) if context.external_data: # Deal with form encoding + external data @@ -303,7 +318,7 @@ class HttpClient(Client): final_query = self._prep_query(context) fields = {} # Setup additional query parameters and body - if self.form_encode_query_params: + if use_form: body = b"" fields["query"] = final_query fields.update(context.bind_params) @@ -490,14 +505,19 @@ class HttpClient(Client): """ try: body = "" + full_body = "" try: raw_body = get_response_data(response) - body = common.format_error(raw_body.decode(errors="backslashreplace")).strip() + full_body = raw_body.decode(errors="backslashreplace") + body = common.format_error(full_body).strip() except Exception: logger.warning("Failed to read error response body", exc_info=True) + err_code = response.headers.get(ex_header) + code = error_code_from_header(err_code) + name = error_name_from_body(full_body) if self.show_clickhouse_errors else None + if self.show_clickhouse_errors: - err_code = response.headers.get(ex_header) if err_code: err_str = f"Received ClickHouse exception, code: {err_code}" else: @@ -513,7 +533,8 @@ class HttpClient(Client): finally: response.close() - raise OperationalError(err_str) if retried else DatabaseError(err_str) from None + err_type = OperationalError if retried else DatabaseError + raise err_type(err_str, code=code, name=name) from None def _raw_request( self, @@ -532,6 +553,7 @@ class HttpClient(Client): data = data.encode() headers = dict_copy(self.headers, headers) attempts = 0 + auth_retried = False final_params = {} if server_wait: final_params["wait_end_of_query"] = "1" @@ -606,6 +628,17 @@ class HttpClient(Client): if attempts > retries: self._error_handler(response, True) logger.debug("Retrying requests with status code %d", response.status) + elif self._token_provider and not auth_retried and response.headers.get(ex_header) == auth_failed_ex_code: + body = kwargs.get("body") + if retry_body is None and not (body is None or isinstance(body, (bytes, bytearray, str))): + self._error_handler(response) # non-replayable body, surface the auth error instead of retrying + auth_retried = True + self.set_access_token(self._token_provider()) + headers["Authorization"] = self.headers["Authorization"] + if retry_body is not None: + kwargs["body"] = retry_body() + response.close() + logger.debug("Refreshing access token after authentication failure") elif error_handler: error_handler(response) else: @@ -667,11 +700,12 @@ class HttpClient(Client): if use_database and self.database: params["database"] = self.database fields = {} + use_form = use_form_encoding(final_query, bind_params, self.form_encode_query_params) # Setup query body - if external_data and not self.form_encode_query_params and isinstance(final_query, bytes): + if external_data and not use_form and isinstance(final_query, bytes): raise ProgrammingError("Binary query cannot be placed in URL when using External Data; enable form encoding.") # Setup additional query parameters and body - if self.form_encode_query_params: + if use_form: body = b"" fields["query"] = final_query fields.update(bind_params) @@ -739,7 +773,12 @@ class HttpClient(Client): See BaseClient doc_string for this method """ try: - response = self.http.request("GET", f"{self.url}/ping", timeout=3, preload_content=True) + headers = dict_copy(self.headers) + kwargs = {"headers": headers, "timeout": 3, "preload_content": True} + if self.server_host_name: + kwargs["assert_same_host"] = False + headers["Host"] = self.server_host_name + response = self.http.request("GET", f"{self.url}/ping", **kwargs) return 200 <= response.status < 300 except HTTPError: logger.debug("ping failed", exc_info=True) diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/query.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/query.py index c17cea1c485..9338e6ff572 100644 --- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/query.py +++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/query.py @@ -403,7 +403,7 @@ class QueryResult(Closable): self._block_gen = None -comment_re = re.compile(r"(\".*?\"|\'.*?\')|(/\*.*?\*/|(--\s)[^\n]*$)", re.MULTILINE | re.DOTALL) +comment_re = re.compile(r"(\".*?\"|\'.*?\')|(/\*.*?\*/|(--)[^\n]*$)", re.MULTILINE | re.DOTALL) def remove_sql_comments(sql: str) -> str: diff --git a/contrib/python/clickhouse-connect/ya.make b/contrib/python/clickhouse-connect/ya.make index 82a5eb6d189..34bf63b5b78 100644 --- a/contrib/python/clickhouse-connect/ya.make +++ b/contrib/python/clickhouse-connect/ya.make @@ -2,7 +2,7 @@ PY3_LIBRARY() -VERSION(1.1.1) +VERSION(1.3.0) LICENSE(Apache-2.0) diff --git a/contrib/python/cryptography/next/rust/.yandex_meta/build.ym b/contrib/python/cryptography/next/rust/.yandex_meta/build.ym index 27c37abb849..1a9f1137767 100644 --- a/contrib/python/cryptography/next/rust/.yandex_meta/build.ym +++ b/contrib/python/cryptography/next/rust/.yandex_meta/build.ym @@ -26,6 +26,7 @@ _rust _openssl {% endblock %} +{% block rust_version %}nightly-2026-06-04{% endblock %} {% block current_version %}41.0.6{% endblock %} {% block current_url %}https://github.com/pyca/cryptography/archive/refs/tags/{{self.version().strip()}}.tar.gz{% endblock %} @@ -46,17 +47,21 @@ export CC_aarch64_apple_darwin=mac-target-cc-arm export DEP_OPENSSL_INCLUDE={{arcadia_root}}/contrib/libs/openssl/include export PYO3_CROSS_PYTHON_VERSION=3.12 +export RUSTFLAGS="${RUSTFLAGS:-} -Cembed-bitcode=no" rm -rf docs vectors tests cd src/rust {{super()}} +find ${BUILD} -name '*.a' -type f | while read l; do + echo "_rust_eh_personality {{self.mod_name().strip()}}__rust_eh_personality" >> prog + llvm-objcopy --redefine-syms=prog ${l} + rm prog +done {% for p in self.msan_platforms().strip().split('\n') %} export CC_{{p.replace('-', '_')}}=target-cc-msan {% endfor %} -export RUSTFLAGS="-Zsanitizer=memory -Zsanitizer-memory-track-origins=2 -Ctarget-feature=-crt-static" - -export RUSTC_BOOTSTRAP=1 +export RUSTFLAGS="${RUSTFLAGS:-} -Zsanitizer=memory -Zsanitizer-memory-track-origins=2 -Ctarget-feature=-crt-static -Cunsafe-allow-abi-mismatch=sanitizer -Ccodegen-units=1" {% for arch in self.msan_platforms().strip().split('\n') %} LD_PRELOAD=${BIN}/getrandom.so cargo {{self.cargo_flags().strip()}} build {{self.cargo_build_flags().strip()}} --release --target {{arch}} {{self.binary_type().strip()}} --target-dir ${BUILD}-msan \ @@ -69,6 +74,7 @@ find ${BUILD}-msan -name '*.a' -type f | while read l; do llvm-nm --extern-only --defined-only -A ${l} | sed -e 's|.* ||' | grep -v 'PyInit_' | sort | while read s; do echo "${s} {{self.mod_name().strip()}}_${s}" >> prog done + echo "_rust_eh_personality {{self.mod_name().strip()}}__rust_eh_personality" >> prog llvm-objcopy --redefine-syms=prog ${l} rm prog done diff --git a/contrib/python/ydb/py3/.dist-info/METADATA b/contrib/python/ydb/py3/.dist-info/METADATA index 3c0337ed5e0..809298530a7 100644 --- a/contrib/python/ydb/py3/.dist-info/METADATA +++ b/contrib/python/ydb/py3/.dist-info/METADATA @@ -1,6 +1,6 @@ Metadata-Version: 2.1 Name: ydb -Version: 3.28.4 +Version: 3.29.6 Summary: YDB Python SDK Home-page: http://github.com/ydb-platform/ydb-python-sdk Author: Yandex LLC @@ -21,6 +21,8 @@ Requires-Dist: grpcio >=1.42.0 Requires-Dist: packaging Requires-Dist: protobuf <7.0.0,>=3.13.0 Requires-Dist: aiohttp <4 +Provides-Extra: opentelemetry +Requires-Dist: opentelemetry-api >=1.0.0 ; extra == 'opentelemetry' Provides-Extra: yc Requires-Dist: yandexcloud ; extra == 'yc' @@ -31,6 +33,7 @@ YDB Python SDK [](https://ydb-platform.github.io/ydb-python-sdk) [](https://github.com/ydb-platform/ydb-python-sdk/actions/workflows/tests.yaml) [](https://github.com/ydb-platform/ydb-python-sdk/actions/workflows/style.yaml) +[](https://codecov.io/gh/ydb-platform/ydb-python-sdk) Officially supported Python client for YDB. diff --git a/contrib/python/ydb/py3/README.md b/contrib/python/ydb/py3/README.md index 7842bd1d910..eeda8db3b42 100644 --- a/contrib/python/ydb/py3/README.md +++ b/contrib/python/ydb/py3/README.md @@ -5,6 +5,7 @@ YDB Python SDK [](https://ydb-platform.github.io/ydb-python-sdk) [](https://github.com/ydb-platform/ydb-python-sdk/actions/workflows/tests.yaml) [](https://github.com/ydb-platform/ydb-python-sdk/actions/workflows/style.yaml) +[](https://codecov.io/gh/ydb-platform/ydb-python-sdk) Officially supported Python client for YDB. diff --git a/contrib/python/ydb/py3/patches/03-change-inport-protobufs.patch b/contrib/python/ydb/py3/patches/03-change-inport-protobufs.patch index a865f95a382..156c6d7799e 100644 --- a/contrib/python/ydb/py3/patches/03-change-inport-protobufs.patch +++ b/contrib/python/ydb/py3/patches/03-change-inport-protobufs.patch @@ -87,10 +87,11 @@ from ..common.protos import ydb_topic_pb2 --- contrib/python/ydb/py3/ydb/aio/connection.py (index) +++ contrib/python/ydb/py3/ydb/aio/connection.py (working tree) -@@ -27,11 +27,10 @@ from ydb.driver import DriverConfig +@@ -28,12 +28,11 @@ from ydb import issues from ydb.settings import BaseRequestSettings from ydb import issues - + from ydb.opentelemetry.tracing import get_trace_metadata + -# Workaround for good IDE and universal for runtime -if TYPE_CHECKING: - from ydb._grpc.v4 import ydb_topic_v1_pb2_grpc diff --git a/contrib/python/ydb/py3/ya.make b/contrib/python/ydb/py3/ya.make index b75adf12a55..5c39099e85f 100644 --- a/contrib/python/ydb/py3/ya.make +++ b/contrib/python/ydb/py3/ya.make @@ -2,7 +2,7 @@ PY3_LIBRARY() -VERSION(3.28.4) +VERSION(3.29.6) LICENSE(Apache-2.0) @@ -16,6 +16,7 @@ PEERDIR( NO_LINT() NO_CHECK_IMPORTS( + ydb.opentelemetry.plugin ydb.public.api.grpc ydb.public.api.grpc.* ) @@ -104,6 +105,9 @@ PY_SRCS( ydb/oauth2_token_exchange/__init__.py ydb/oauth2_token_exchange/token_exchange.py ydb/oauth2_token_exchange/token_source.py + ydb/opentelemetry/__init__.py + ydb/opentelemetry/plugin.py + ydb/opentelemetry/tracing.py ydb/operation.py ydb/pool.py ydb/query/__init__.py diff --git a/contrib/python/ydb/py3/ydb/_topic_reader/topic_reader_asyncio.py b/contrib/python/ydb/py3/ydb/_topic_reader/topic_reader_asyncio.py index f4e5a4f88a5..ee988a1ec3b 100644 --- a/contrib/python/ydb/py3/ydb/_topic_reader/topic_reader_asyncio.py +++ b/contrib/python/ydb/py3/ydb/_topic_reader/topic_reader_asyncio.py @@ -217,6 +217,7 @@ class ReaderReconnector: _stream_reader: Optional["ReaderStream"] _first_error: asyncio.Future[YdbError] _tx_to_batches_map: Dict[str, typing.List[datatypes.PublicBatch]] + _closed: bool def __init__( self, @@ -233,6 +234,7 @@ class ReaderReconnector: self._state_changed = asyncio.Event() self._stream_reader = None + self._closed = False self._background_tasks.add(asyncio.create_task(self._connection_loop())) self._first_error = asyncio.get_running_loop().create_future() @@ -241,6 +243,8 @@ class ReaderReconnector: async def _connection_loop(self): attempt = 0 while True: + if self._closed: + return try: logger.debug("reader %s connect attempt %s", self._id, attempt) self._stream_reader = await ReaderStream.create(self._id, self._driver, self._settings) @@ -266,8 +270,12 @@ class ReaderReconnector: # noinspection PyBroadException try: await self._stream_reader.close(flush=False) - except BaseException: - # supress any error on close stream reader + except asyncio.CancelledError: + # propagate cancellation (e.g. from reader.close()) so the loop stops + # instead of swallowing it and reconnecting into a zombie stream + raise + except Exception: + # suppress any error on close stream reader pass async def wait_message(self): @@ -323,14 +331,44 @@ class ReaderReconnector: tx._add_callback(TxEvent.AFTER_COMMIT, self._handle_after_tx_commit, self._loop) tx._add_callback(TxEvent.AFTER_ROLLBACK, self._handle_after_tx_rollback, self._loop) + def _batch_partition_session_expired(self, batch: datatypes.PublicBatch) -> bool: + # A batch is expired if the reader reconnected after it was received: its partition + # session no longer belongs to the current stream. Mirrors the guard in + # ReaderStream.commit() for the non-transactional commit path. + stream = self._stream_reader + partition_session = batch._partition_session + return ( + stream is None + or partition_session.reader_stream_id != stream._id + or partition_session.id not in stream._partition_sessions + ) + async def _commit_batches_with_tx(self, tx: "BaseQueryTxContext"): tx_id = tx.tx_id if tx_id is None: raise TopicReaderError("Transaction ID is None") + + batches = self._tx_to_batches_map[tx_id] + + if any(self._batch_partition_session_expired(batch) for batch in batches): + # The reader reconnected between receive_batch_with_tx() and tx.commit(), so + # these offsets belong to a partition session that no longer exists. Committing + # them would send a stale/gapped range (server "Gap", issue_code 2011) while the + # client believes the commit succeeded. Fail the tx instead (retriable) without + # sending the request; the AFTER_COMMIT handler then reconnects to reset the + # read-ahead state, and the pool re-reads from the committed offset. + err = issues.ClientInternalError( + "Topic reader partition session expired before tx commit; " + "offsets were not committed, the transaction will be retried" + ) + tx._set_external_error(err) + del self._tx_to_batches_map[tx_id] + return + grouped_batches: Dict[str, Dict[int, typing.List[datatypes.PublicBatch]]] = defaultdict( lambda: defaultdict(list) ) - for batch in self._tx_to_batches_map[tx_id]: + for batch in batches: grouped_batches[batch._partition_session.topic_path][batch._partition_session.partition_id].append(batch) consumer = self._settings.consumer @@ -401,8 +439,16 @@ class ReaderReconnector: async def close(self, flush: bool): logger.debug("reader reconnector %s close", self._id) + # Mark closed so the connection loop won't start a new stream, then close the + # current stream with the requested flush before cancelling the loop. On a normal + # close this flushes pending commits; cancelling the loop first would let it close + # the stream with flush=False instead and skip the flush. + self._closed = True if self._stream_reader: await self._stream_reader.close(flush) + # Wake any pending wait_message() waiter (e.g. a concurrent receive) so it doesn't + # hang if the loop was reconnecting when close() cancelled it. + self._set_first_error(TopicReaderStreamClosedError()) for task in self._background_tasks: task.cancel() @@ -469,6 +515,11 @@ class ReaderStream: self._id = ReaderStream._static_id_counter.inc_and_get() self._reader_reconnector_id = reader_reconnector_id self._session_id = "not initialized" + self._log_prefix = "reader %s stream %s session=%s" % ( + self._reader_reconnector_id, + self._id, + self._session_id, + ) self._stream = None self._started = False self._background_tasks = set() @@ -503,17 +554,28 @@ class ReaderStream: settings: topic_reader.PublicReaderSettings, ) -> "ReaderStream": stream = GrpcWrapperAsyncIO(StreamReadMessage.FromServer.from_proto) + reader = None + try: + await stream.start(driver, _apis.TopicService.Stub, _apis.TopicService.StreamRead) - await stream.start(driver, _apis.TopicService.Stub, _apis.TopicService.StreamRead) - - creds = driver._credentials - reader = ReaderStream( - reader_reconnector_id, - settings, - get_token_function=creds.get_auth_token if creds else None, - ) - await reader._start(stream, settings._init_message()) - logger.debug("reader stream %s started session=%s", reader._id, reader._session_id) + creds = driver._credentials + reader = ReaderStream( + reader_reconnector_id, + settings, + get_token_function=creds.get_auth_token if creds else None, + ) + await reader._start(stream, settings._init_message()) + except BaseException: + # If create() is interrupted (e.g. reader.close() cancels the connection loop + # mid-reconnect) the in-flight stream is not yet assigned to the reconnector, so + # its finally cannot reach it. Close it here to avoid a zombie gRPC read session + # that keeps holding the consumer's partition on the server. + if reader is not None: + await reader.close(flush=False) + else: + stream.close() + raise + logger.debug("%s started", reader._log_prefix) return reader async def _start(self, stream: IGrpcWrapperAsyncIO, init_message: StreamReadMessage.InitRequest): @@ -522,7 +584,7 @@ class ReaderStream: self._started = True self._stream = stream - logger.debug("reader stream %s send init request", self._id) + logger.debug("%s send init request", self._log_prefix) stream.write(StreamReadMessage.FromClient(client_message=init_message)) try: @@ -534,7 +596,12 @@ class ReaderStream: if isinstance(init_response.server_message, StreamReadMessage.InitResponse): self._session_id = init_response.server_message.session_id - logger.debug("reader stream %s initialized session=%s", self._id, self._session_id) + self._log_prefix = "reader %s stream %s session=%s" % ( + self._reader_reconnector_id, + self._id, + self._session_id, + ) + logger.debug("%s initialized", self._log_prefix) else: raise TopicReaderError("Unexpected message after InitRequest: %s" % init_response) @@ -668,7 +735,7 @@ class ReaderStream: async def _read_messages_loop(self): try: - logger.debug("reader stream %s start read loop", self._id) + logger.debug("%s start read loop", self._log_prefix) self._stream.write( StreamReadMessage.FromClient( client_message=StreamReadMessage.ReadRequest( @@ -682,7 +749,7 @@ class ReaderStream: _process_response(message.server_status) if isinstance(message.server_message, StreamReadMessage.ReadResponse): - logger.debug("reader stream %s read %s bytes", self._id, message.server_message.bytes_size) + logger.debug("%s read %s bytes", self._log_prefix, message.server_message.bytes_size) self._on_read_response(message.server_message) elif isinstance(message.server_message, StreamReadMessage.CommitOffsetResponse): @@ -693,8 +760,8 @@ class ReaderStream: StreamReadMessage.StartPartitionSessionRequest, ): logger.debug( - "reader stream %s start partition %s", - self._id, + "%s start partition %s", + self._log_prefix, message.server_message.partition_session.partition_session_id, ) await self._on_start_partition_session(message.server_message) @@ -704,8 +771,8 @@ class ReaderStream: StreamReadMessage.StopPartitionSessionRequest, ): logger.debug( - "reader stream %s stop partition %s", - self._id, + "%s stop partition %s", + self._log_prefix, message.server_message.partition_session_id, ) self._on_partition_session_stop(message.server_message) @@ -715,8 +782,8 @@ class ReaderStream: StreamReadMessage.EndPartitionSession, ): logger.debug( - "reader stream %s end partition %s", - self._id, + "%s end partition %s", + self._log_prefix, message.server_message.partition_session_id, ) self._on_end_partition_session(message.server_message) @@ -733,12 +800,12 @@ class ReaderStream: self._state_changed.set() except asyncio.CancelledError as e: - logger.debug("reader stream %s error: %s", self._id, e) + logger.debug("%s error: %s", self._log_prefix, e) if not self._closed: self._set_first_error(issues.ConnectionLost("gRPC stream cancelled")) raise except Exception as e: - logger.debug("reader stream %s error: %s", self._id, e) + logger.debug("%s error: %s", self._log_prefix, e) self._set_first_error(e) return @@ -908,7 +975,7 @@ class ReaderStream: async def _decode_batches_loop(self): while True: batch = await self._batches_to_decode.get() - logger.debug("reader stream %s decode batch %s messages", self._id, len(batch.messages)) + logger.debug("%s decode batch %s messages", self._log_prefix, len(batch.messages)) await self._decode_batch_inplace(batch) self._add_batch_to_queue(batch) self._state_changed.set() @@ -918,8 +985,8 @@ class ReaderStream: if part_sess_id in self._message_batches: self._message_batches[part_sess_id]._extend(batch) logger.debug( - "reader stream %s extend batch partition=%s size=%s", - self._id, + "%s extend batch partition=%s size=%s", + self._log_prefix, part_sess_id, len(batch.messages), ) @@ -927,8 +994,8 @@ class ReaderStream: self._message_batches[part_sess_id] = batch logger.debug( - "reader stream %s new batch partition=%s size=%s", - self._id, + "%s new batch partition=%s size=%s", + self._log_prefix, part_sess_id, len(batch.messages), ) @@ -979,7 +1046,7 @@ class ReaderStream: return self._closed = True - logger.debug("reader stream %s close", self._id) + logger.debug("%s close", self._log_prefix) if flush: await self.flush() @@ -999,4 +1066,4 @@ class ReaderStream: if self._background_tasks: await asyncio.wait(self._background_tasks) - logger.debug("reader stream %s was closed", self._id) + logger.debug("%s was closed", self._log_prefix) diff --git a/contrib/python/ydb/py3/ydb/_topic_writer/topic_writer_asyncio.py b/contrib/python/ydb/py3/ydb/_topic_writer/topic_writer_asyncio.py index 5e55917d830..ff9c59f7989 100644 --- a/contrib/python/ydb/py3/ydb/_topic_writer/topic_writer_asyncio.py +++ b/contrib/python/ydb/py3/ydb/_topic_writer/topic_writer_asyncio.py @@ -833,7 +833,8 @@ class WriterAsyncIOStream: self._update_token_task.cancel() await asyncio.wait([self._update_token_task]) - self._stream.close() + if getattr(self, "_stream", None) is not None: + self._stream.close() logger.debug("writer stream %s was closed", self._id) @staticmethod @@ -845,15 +846,27 @@ class WriterAsyncIOStream: ) -> "WriterAsyncIOStream": stream = GrpcWrapperAsyncIO(StreamWriteMessage.FromServer.from_proto) - await stream.start(driver, _apis.TopicService.Stub, _apis.TopicService.StreamWrite) + writer = None + try: + await stream.start(driver, _apis.TopicService.Stub, _apis.TopicService.StreamWrite) - creds = driver._credentials - writer = WriterAsyncIOStream( - update_token_interval=update_token_interval, - get_token_function=creds.get_auth_token if creds else lambda: "", - tx_identity=tx_identity, - ) - await writer._start(stream, init_request) + creds = driver._credentials + writer = WriterAsyncIOStream( + update_token_interval=update_token_interval, + get_token_function=creds.get_auth_token if creds else lambda: "", + tx_identity=tx_identity, + ) + await writer._start(stream, init_request) + except BaseException: + # If create() fails after stream.start() (e.g. the connection is lost while + # waiting for the init response), the stream is not yet handed to the + # reconnector, so its finally cannot reach it. Close it here to avoid a + # stranded gRPC consumption thread that blocks forever on the request queue. + if writer is not None: + await writer.close() + else: + stream.close() + raise logger.debug( "writer stream %s started seqno=%s", writer._id, @@ -875,6 +888,10 @@ class WriterAsyncIOStream: raise Exception("Unknown message while read writer answers: %s" % item) async def _start(self, stream: IGrpcWrapperAsyncIO, init_message: StreamWriteMessage.InitRequest): + # Assign before the init handshake (mirrors ReaderStream._start) so that if _start + # fails mid-handshake, close() can still reach the stream and release its gRPC thread. + self._stream = stream + logger.debug("writer stream %s send init request", self._id) stream.write(StreamWriteMessage.FromClient(init_message)) @@ -895,8 +912,6 @@ class WriterAsyncIOStream: self.last_seqno, ) - self._stream = stream - if self._update_token_interval is not None: self._update_token_event.set() self._update_token_task = asyncio.create_task(self._update_token_loop()) diff --git a/contrib/python/ydb/py3/ydb/aio/connection.py b/contrib/python/ydb/py3/ydb/aio/connection.py index 3cb6290e1d4..02320bf9afc 100644 --- a/contrib/python/ydb/py3/ydb/aio/connection.py +++ b/contrib/python/ydb/py3/ydb/aio/connection.py @@ -26,6 +26,7 @@ from ydb.connection import ( from ydb.driver import DriverConfig from ydb.settings import BaseRequestSettings from ydb import issues +from ydb.opentelemetry.tracing import get_trace_metadata try: from ydb.public.api.grpc import ydb_topic_v1_pb2_grpc @@ -70,6 +71,9 @@ async def _construct_metadata( metadata.append((YDB_REQUEST_TYPE_HEADER, settings.request_type)) metadata.append(_utilities.x_ydb_sdk_build_info_header(getattr(driver_config, "_additional_sdk_headers", ()))) + + metadata.extend(get_trace_metadata()) + return metadata @@ -148,6 +152,9 @@ class Connection: "closing", "endpoint_key", "node_id", + "peer_address", + "peer_port", + "peer_location", ) def __init__( @@ -156,10 +163,12 @@ class Connection: driver_config: Optional[DriverConfig] = None, endpoint_options: Optional[EndpointOptions] = None, ) -> None: - global _stubs_list self.endpoint = endpoint self.endpoint_key = EndpointKey(self.endpoint, getattr(endpoint_options, "node_id", None)) self.node_id = getattr(endpoint_options, "node_id", None) + self.peer_address = getattr(endpoint_options, "address", None) + self.peer_port = getattr(endpoint_options, "port", None) + self.peer_location = getattr(endpoint_options, "location", None) self._channel = channel_factory(self.endpoint, driver_config, grpc.aio, endpoint_options=endpoint_options) self._driver_config = driver_config @@ -248,8 +257,12 @@ class Connection: :param grace: :return: None """ - if hasattr(self, "_channel") and hasattr(self._channel, "close"): - await self._channel.close(grace) + channel = getattr(self, "_channel", None) + if channel is not None and hasattr(channel, "close"): + await channel.close(grace) + + self._stub_instances.clear() + self._channel = None def add_cleanup_callback(self, callback: Callable[["Connection"], None]) -> None: self._cleanup_callbacks.append(callback) diff --git a/contrib/python/ydb/py3/ydb/aio/driver.py b/contrib/python/ydb/py3/ydb/aio/driver.py index 405e5fcbf37..88221e947db 100644 --- a/contrib/python/ydb/py3/ydb/aio/driver.py +++ b/contrib/python/ydb/py3/ydb/aio/driver.py @@ -69,6 +69,7 @@ class Driver(pool.ConnectionPool): root_certificates, credentials, config_class=DriverConfig, + **kwargs, ) super(Driver, self).__init__(config) diff --git a/contrib/python/ydb/py3/ydb/aio/pool.py b/contrib/python/ydb/py3/ydb/aio/pool.py index fe7091332ce..4d952086c55 100644 --- a/contrib/python/ydb/py3/ydb/aio/pool.py +++ b/contrib/python/ydb/py3/ydb/aio/pool.py @@ -6,6 +6,7 @@ import random from typing import Any, Callable, Optional, Tuple, TYPE_CHECKING from ydb import issues +from ydb.opentelemetry.tracing import SpanName, create_ydb_span from ydb.pool import ConnectionsCache as _ConnectionsCache, IConnectionPool from .connection import Connection, EndpointKey @@ -247,17 +248,36 @@ class ConnectionPool(IConnectionPool): self._store = ConnectionsCache(driver_config.use_all_nodes) self._grpc_init = Connection(self._driver_config.endpoint, self._driver_config) self._stopped = False + self._stopping = False self._discovery: Optional[Discovery] = None self._discovery_task: "asyncio.Task[None]" if driver_config.disable_discovery: # If discovery is disabled, just add the initial endpoint to the store async def init_connection() -> None: - ready_connection = Connection(self._driver_config.endpoint, self._driver_config) - await ready_connection.connection_ready( - ready_timeout=getattr(self._driver_config, "discovery_request_timeout", 10) - ) - self._store.add(ready_connection) + ready_timeout = getattr(self._driver_config, "discovery_request_timeout", 10) + while not self._stopping: + ready_connection = Connection(self._driver_config.endpoint, self._driver_config) + try: + await ready_connection.connection_ready(ready_timeout=ready_timeout) + except asyncio.CancelledError: + try: + await ready_connection.close() + except Exception: + logger.debug("Failed to close cancelled initial connection", exc_info=True) + raise + except Exception: + logger.debug("Initial connection attempt failed", exc_info=True) + try: + await ready_connection.close() + except Exception: + logger.debug("Failed to close unsuccessful initial connection", exc_info=True) + if not self._stopping: + await asyncio.sleep(1) + continue + + self._store.add(ready_connection) + return # Create and schedule the task to initialize the connection self._discovery_task = asyncio.get_event_loop().create_task(init_connection()) @@ -267,6 +287,7 @@ class ConnectionPool(IConnectionPool): self._discovery_task = asyncio.get_event_loop().create_task(self._discovery.run()) async def stop(self, timeout: int = 10) -> None: # type: ignore[override] # async override of sync method + self._stopping = True if self._discovery: self._discovery.stop() await self._grpc_init.close() @@ -274,6 +295,12 @@ class ConnectionPool(IConnectionPool): await asyncio.wait_for(self._discovery_task, timeout=timeout) except asyncio.TimeoutError: self._discovery_task.cancel() + try: + await self._discovery_task + except asyncio.CancelledError: + pass + if self._discovery is None: + await self._store.cleanup() self._stopped = True def _on_disconnected(self, connection: Connection) -> Callable[[], Any]: @@ -285,7 +312,8 @@ class ConnectionPool(IConnectionPool): return __wrapper__ async def wait(self, timeout: Optional[float] = 7.0, fail_fast: bool = False) -> None: # type: ignore[override] # async override of sync method - await self._store.get(fast_fail=fail_fast, wait_timeout=timeout if timeout is not None else 7.0) + with create_ydb_span(SpanName.DRIVER_INITIALIZE, self._driver_config, kind="internal").attach_context(): + await self._store.get(fast_fail=fail_fast, wait_timeout=timeout if timeout is not None else 7.0) def discovery_debug_details(self) -> str: if self._discovery: diff --git a/contrib/python/ydb/py3/ydb/aio/query/base.py b/contrib/python/ydb/py3/ydb/aio/query/base.py index a1d6e5d74fc..6c13dfd400a 100644 --- a/contrib/python/ydb/py3/ydb/aio/query/base.py +++ b/contrib/python/ydb/py3/ydb/aio/query/base.py @@ -2,9 +2,12 @@ from .. import _utilities class AsyncResponseContextIterator(_utilities.AsyncResponseIterator): - def __init__(self, it, wrapper, on_error=None): + """Async ExecuteQuery result stream.""" + + def __init__(self, it, wrapper, on_error=None, on_finish=None): super().__init__(it, wrapper) self._on_error = on_error + self._on_finish = on_finish async def __aenter__(self) -> "AsyncResponseContextIterator": return self @@ -15,6 +18,7 @@ class AsyncResponseContextIterator(_utilities.AsyncResponseIterator): except StopAsyncIteration: # Normal stream termination is not an error and must not invalidate # the session. + self._call_on_finish() raise except BaseException as e: # BaseException (not Exception) because asyncio.CancelledError @@ -25,8 +29,17 @@ class AsyncResponseContextIterator(_utilities.AsyncResponseIterator): # reply with SessionBusy. if self._on_error: self._on_error(e) + self._call_on_finish(e) raise + def _call_on_finish(self, exception=None): + if self._on_finish is not None: + self._on_finish(exception) + self._on_finish = None + + def __del__(self): + self._call_on_finish() + async def __aexit__(self, exc_type, exc_val, exc_tb): # To close stream on YDB it is necessary to scroll through it to the end. # Errors that happen during the cleanup drain have already been reported @@ -39,3 +52,4 @@ class AsyncResponseContextIterator(_utilities.AsyncResponseIterator): pass except BaseException: pass + self._call_on_finish() diff --git a/contrib/python/ydb/py3/ydb/aio/query/session.py b/contrib/python/ydb/py3/ydb/aio/query/session.py index a565b266e82..b776b638268 100644 --- a/contrib/python/ydb/py3/ydb/aio/query/session.py +++ b/contrib/python/ydb/py3/ydb/aio/query/session.py @@ -19,6 +19,7 @@ from ..._grpc.grpcwrapper import ydb_query_public_types as _ydb_query_public from ...query import base from ...query.session import BaseQuerySession +from ...opentelemetry.tracing import SpanName, create_ydb_span, set_peer_attributes, span_finish_callback from ..._constants import DEFAULT_INITIAL_RESPONSE_TIMEOUT @@ -105,8 +106,10 @@ class QuerySession(BaseQuerySession["AsyncDriver"]): if self._closed: raise RuntimeError("Session is already closed") - await self._create_call(settings=settings) - await self._attach() + with create_ydb_span(SpanName.CREATE_SESSION, self._driver_config).attach_context() as span: + await self._create_call(settings=settings) + set_peer_attributes(span, self._peer) + await self._attach() return self @@ -159,20 +162,27 @@ class QuerySession(BaseQuerySession["AsyncDriver"]): """ self._check_session_ready_to_use() - stream_it = await self._execute_call( - query=query, - parameters=parameters, - commit_tx=True, - syntax=syntax, - exec_mode=exec_mode, - stats_mode=stats_mode, - schema_inclusion_mode=schema_inclusion_mode, - result_set_format=result_set_format, - arrow_format_settings=arrow_format_settings, - concurrent_result_sets=concurrent_result_sets, - settings=settings, + span = create_ydb_span( + SpanName.EXECUTE_QUERY, + self._driver_config, + node_id=self._node_id, + peer=self._peer, ) + with span.attach_context(end_on_exit=False): + stream_it = await self._execute_call( + query=query, + parameters=parameters, + commit_tx=True, + syntax=syntax, + exec_mode=exec_mode, + stats_mode=stats_mode, + schema_inclusion_mode=schema_inclusion_mode, + result_set_format=result_set_format, + arrow_format_settings=arrow_format_settings, + concurrent_result_sets=concurrent_result_sets, + settings=settings, + ) return AsyncResponseContextIterator( it=stream_it, wrapper=lambda resp: base.wrap_execute_query_response( @@ -182,6 +192,7 @@ class QuerySession(BaseQuerySession["AsyncDriver"]): settings=self._settings, ), on_error=self._on_execute_stream_error, + on_finish=span_finish_callback(span), ) async def explain( diff --git a/contrib/python/ydb/py3/ydb/aio/query/transaction.py b/contrib/python/ydb/py3/ydb/aio/query/transaction.py index c31d79fb088..a05d91f2349 100644 --- a/contrib/python/ydb/py3/ydb/aio/query/transaction.py +++ b/contrib/python/ydb/py3/ydb/aio/query/transaction.py @@ -12,6 +12,7 @@ from ...query.transaction import ( BaseQueryTxContext, QueryTxStateEnum, ) +from ...opentelemetry.tracing import SpanName, create_ydb_span, span_finish_callback if TYPE_CHECKING: from .session import QuerySession @@ -87,7 +88,13 @@ class QueryTxContext(BaseQueryTxContext["AsyncDriver"]): :return: None or exception if begin is failed """ - await self._begin_call(settings) + with create_ydb_span( + SpanName.BEGIN_TRANSACTION, + self._driver_config, + node_id=self.session.node_id, + peer=getattr(self.session, "_peer", None), + ).attach_context(): + await self._begin_call(settings) return self async def commit(self, settings: Optional[BaseRequestSettings] = None) -> None: @@ -109,13 +116,19 @@ class QueryTxContext(BaseQueryTxContext["AsyncDriver"]): await self._ensure_prev_stream_finished() - try: - await self._execute_callbacks_async(base.TxEvent.BEFORE_COMMIT) - await self._commit_call(settings) - await self._execute_callbacks_async(base.TxEvent.AFTER_COMMIT, exc=None) - except BaseException as e: - await self._execute_callbacks_async(base.TxEvent.AFTER_COMMIT, exc=e) - raise e + with create_ydb_span( + SpanName.COMMIT, + self._driver_config, + node_id=self.session.node_id, + peer=getattr(self.session, "_peer", None), + ).attach_context(): + try: + await self._execute_callbacks_async(base.TxEvent.BEFORE_COMMIT) + await self._commit_call(settings) + await self._execute_callbacks_async(base.TxEvent.AFTER_COMMIT, exc=None) + except BaseException as e: + await self._execute_callbacks_async(base.TxEvent.AFTER_COMMIT, exc=e) + raise e async def rollback(self, settings: Optional[BaseRequestSettings] = None) -> None: """Calls rollback on a transaction if it is open otherwise is no-op. If transaction execution @@ -136,13 +149,19 @@ class QueryTxContext(BaseQueryTxContext["AsyncDriver"]): await self._ensure_prev_stream_finished() - try: - await self._execute_callbacks_async(base.TxEvent.BEFORE_ROLLBACK) - await self._rollback_call(settings) - await self._execute_callbacks_async(base.TxEvent.AFTER_ROLLBACK, exc=None) - except BaseException as e: - await self._execute_callbacks_async(base.TxEvent.AFTER_ROLLBACK, exc=e) - raise e + with create_ydb_span( + SpanName.ROLLBACK, + self._driver_config, + node_id=self.session.node_id, + peer=getattr(self.session, "_peer", None), + ).attach_context(): + try: + await self._execute_callbacks_async(base.TxEvent.BEFORE_ROLLBACK) + await self._rollback_call(settings) + await self._execute_callbacks_async(base.TxEvent.AFTER_ROLLBACK, exc=None) + except BaseException as e: + await self._execute_callbacks_async(base.TxEvent.AFTER_ROLLBACK, exc=e) + raise e async def execute( self, @@ -190,20 +209,27 @@ class QueryTxContext(BaseQueryTxContext["AsyncDriver"]): """ await self._ensure_prev_stream_finished() - stream_it = await self._execute_call( - query=query, - parameters=parameters, - commit_tx=commit_tx, - syntax=syntax, - exec_mode=exec_mode, - stats_mode=stats_mode, - schema_inclusion_mode=schema_inclusion_mode, - result_set_format=result_set_format, - arrow_format_settings=arrow_format_settings, - concurrent_result_sets=concurrent_result_sets, - settings=settings, + span = create_ydb_span( + SpanName.EXECUTE_QUERY, + self._driver_config, + node_id=self.session.node_id, + peer=getattr(self.session, "_peer", None), ) + with span.attach_context(end_on_exit=False): + stream_it = await self._execute_call( + query=query, + parameters=parameters, + commit_tx=commit_tx, + syntax=syntax, + exec_mode=exec_mode, + stats_mode=stats_mode, + schema_inclusion_mode=schema_inclusion_mode, + result_set_format=result_set_format, + arrow_format_settings=arrow_format_settings, + concurrent_result_sets=concurrent_result_sets, + settings=settings, + ) self._prev_stream = AsyncResponseContextIterator( it=stream_it, wrapper=lambda resp: base.wrap_execute_query_response( @@ -215,5 +241,6 @@ class QueryTxContext(BaseQueryTxContext["AsyncDriver"]): settings=self.session._settings, ), on_error=self.session._on_execute_stream_error, + on_finish=span_finish_callback(span), ) return self._prev_stream diff --git a/contrib/python/ydb/py3/ydb/aio/table.py b/contrib/python/ydb/py3/ydb/aio/table.py index 0d14ba2f11b..8d5e02c180d 100644 --- a/contrib/python/ydb/py3/ydb/aio/table.py +++ b/contrib/python/ydb/py3/ydb/aio/table.py @@ -462,7 +462,8 @@ async def retry_operation(callee, retry_settings=None, *args, **kwargs): # pyli opt_generator = ydb.retry_operation_impl(callee, retry_settings, *args, **kwargs) for next_opt in opt_generator: if isinstance(next_opt, ydb.YdbRetryOperationSleepOpt): - await asyncio.sleep(next_opt.timeout) + if next_opt.timeout > 0: + await asyncio.sleep(next_opt.timeout) else: try: return await next_opt.result diff --git a/contrib/python/ydb/py3/ydb/connection.py b/contrib/python/ydb/py3/ydb/connection.py index 98fbd5aab31..7e5ad3675fe 100644 --- a/contrib/python/ydb/py3/ydb/connection.py +++ b/contrib/python/ydb/py3/ydb/connection.py @@ -24,6 +24,7 @@ from google.protobuf import text_format import grpc from . import issues, _apis, _utilities from . import default_pem +from .opentelemetry.tracing import get_trace_metadata _stubs_list = ( _apis.TableService.Stub, @@ -179,6 +180,9 @@ def _construct_metadata(driver_config, settings): metadata.extend(getattr(settings, "headers", [])) metadata.append(_utilities.x_ydb_sdk_build_info_header(getattr(driver_config, "_additional_sdk_headers", ()))) + + metadata.extend(get_trace_metadata()) + return metadata @@ -194,11 +198,14 @@ def _get_request_timeout(settings): class EndpointOptions(object): - __slots__ = ("ssl_target_name_override", "node_id") + __slots__ = ("ssl_target_name_override", "node_id", "address", "port", "location") - def __init__(self, ssl_target_name_override=None, node_id=None): + def __init__(self, ssl_target_name_override=None, node_id=None, address=None, port=None, location=None): self.ssl_target_name_override = ssl_target_name_override self.node_id = node_id + self.address = address + self.port = port + self.location = location def _construct_channel_options(driver_config, endpoint_options=None): @@ -405,6 +412,9 @@ class Connection(object): "closing", "endpoint_key", "node_id", + "peer_address", + "peer_port", + "peer_location", ) def __init__( @@ -419,9 +429,11 @@ class Connection(object): discovered by the YDB endpoint discovery mechanism :param driver_config: A driver config instance to be used for RPC call interception """ - global _stubs_list self.endpoint = endpoint self.node_id = getattr(endpoint_options, "node_id", None) + self.peer_address = getattr(endpoint_options, "address", None) + self.peer_port = getattr(endpoint_options, "port", None) + self.peer_location = getattr(endpoint_options, "location", None) self.endpoint_key = EndpointKey(endpoint, getattr(endpoint_options, "node_id", None)) self._channel = channel_factory(self.endpoint, driver_config, endpoint_options=endpoint_options) self._driver_config = driver_config @@ -590,8 +602,12 @@ class Connection(object): self.destroy() def destroy(self): - if hasattr(self, "_channel") and hasattr(self._channel, "close"): - self._channel.close() + channel = getattr(self, "_channel", None) + if channel is not None and hasattr(channel, "close"): + channel.close() + + self._stub_instances.clear() + self._channel = None def ready_future(self): """ diff --git a/contrib/python/ydb/py3/ydb/convert.py b/contrib/python/ydb/py3/ydb/convert.py index b54e8bddc8d..89f28ecc597 100644 --- a/contrib/python/ydb/py3/ydb/convert.py +++ b/contrib/python/ydb/py3/ydb/convert.py @@ -24,11 +24,21 @@ _initialize() class _DotDict(dict): + # A lazy __dict__ is declared on purpose: it is not materialized until a row + # is written to (or its __dict__ is introspected), so untouched read-only + # rows avoid the per-instance dict overhead while callers can still attach + # their own attributes (ORM-style row.extra = ...), which materializes the + # dict for that row alone. + __slots__ = ("__dict__",) + def __init__(self, *args, **kwargs): super(_DotDict, self).__init__(*args, **kwargs) def __getattr__(self, item): - return self[item] + try: + return self[item] + except KeyError: + raise AttributeError(item) from None def _is_decimal_signed(hi_value): @@ -83,7 +93,7 @@ def _pb_to_dict(type_pb, value_pb, table_client_settings): class _Struct(_DotDict): - pass + __slots__ = () def _pb_to_struct(type_pb, value_pb, table_client_settings): @@ -351,6 +361,16 @@ def _unwrap_optionality(column): return _to_native_map.get(current_type), c_type +def _detach_columns(columns): + # The columns container references the source protobuf message, which keeps + # the whole arena (including already-parsed message.rows) alive for as long + # as the result set is held. Copy the schema into a standalone message so the + # source arena can be released right after conversion. + holder = _apis.ydb_value.ResultSet() + holder.columns.extend(columns) + return holder.columns + + class _ResultSet(object): __slots__ = ("columns", "rows", "truncated", "snapshot", "index", "format", "arrow_format_meta", "data") @@ -375,12 +395,17 @@ class _ResultSet(object): for column in message.columns: column_parsers.append(_unwrap_optionality(column)) + # Names are the only per-row metadata we need. Storing this shared tuple + # (instead of the proto columns) keeps rows detached from the source + # protobuf arena, so it can be freed once conversion is done. + column_names = tuple(column.name for column in message.columns) + for row_proto in message.rows: - row = _Row(message.columns) - for column, value, column_info in zip(message.columns, row_proto.items, column_parsers): + row = _Row(column_names) + for name, value, column_info in zip(column_names, row_proto.items, column_parsers): v_type = value.WhichOneof("value") if v_type == "null_flag_value": - row[column.name] = None + row[name] = None continue while v_type == "nested_value": @@ -388,7 +413,7 @@ class _ResultSet(object): v_type = value.WhichOneof("value") column_parser, unwrapped_type = column_info - row[column.name] = column_parser(unwrapped_type, value, table_client_settings) + row[name] = column_parser(unwrapped_type, value, table_client_settings) rows.append(row) from ydb.query import QueryResultSetFormat, ArrowFormatMeta @@ -401,12 +426,17 @@ class _ResultSet(object): data = message.data if message.data else None - return cls(message.columns, rows, message.truncated, snapshot, index, result_format, arrow_meta, data) + return cls( + _detach_columns(message.columns), rows, message.truncated, snapshot, index, result_format, arrow_meta, data + ) @classmethod def lazy_from_message(cls, message, table_client_settings=None, snapshot=None): from ydb.query import QueryResultSetFormat, ArrowFormatMeta + # No _detach_columns here on purpose: _LazyRows defers parsing and keeps + # message.rows, so the source arena is pinned regardless — detaching the + # schema would only add a copy without freeing anything. rows = _LazyRows(message.rows, table_client_settings, message.columns) result_format = message.format if message.format else QueryResultSetFormat.VALUE @@ -422,15 +452,17 @@ ResultSet = _ResultSet class _Row(_DotDict): + __slots__ = ("_columns",) + def __init__(self, columns): super(_Row, self).__init__() self._columns = columns def __getitem__(self, key): if isinstance(key, int): - return self[self._columns[key].name] + return self[self._columns[key]] elif isinstance(key, slice): - return tuple(map(lambda x: self[x.name], self._columns[key])) + return tuple(self[name] for name in self._columns[key]) else: return super(_Row, self).__getitem__(key) @@ -455,10 +487,11 @@ class _LazyRowItem: class _LazyRow(_DotDict): + __slots__ = ("_columns",) + def __init__(self, columns, proto_row, table_client_settings, parsers): super(_LazyRow, self).__init__() self._columns = columns - self._table_client_settings = table_client_settings for i, (column, row_item) in enumerate(zip(self._columns, proto_row.items)): super(_LazyRow, self).__setitem__( column.name, diff --git a/contrib/python/ydb/py3/ydb/opentelemetry/__init__.py b/contrib/python/ydb/py3/ydb/opentelemetry/__init__.py new file mode 100644 index 00000000000..fc058d0d866 --- /dev/null +++ b/contrib/python/ydb/py3/ydb/opentelemetry/__init__.py @@ -0,0 +1,36 @@ +"""Public OpenTelemetry entrypoints for YDB.""" + + +def enable_tracing(tracer=None): + """Enable OpenTelemetry trace context propagation and span creation for all YDB gRPC calls. + + This call is **idempotent**: if tracing is already enabled, later calls do nothing + (including passing a different ``tracer``). Call :func:`disable_tracing` first to + reconfigure or turn instrumentation off. + + Args: + tracer: Optional OTel tracer to use. If not provided, the default tracer named + ``ydb.sdk`` from the global tracer provider will be used. + """ + try: + from ydb.opentelemetry.plugin import _enable_tracing + except ImportError: + raise ImportError( + "OpenTelemetry packages are required for tracing support. " + "Install them with: pip install ydb[opentelemetry]" + ) from None + + _enable_tracing(tracer) + + +def disable_tracing(): + """Disable YDB OpenTelemetry hooks and allow :func:`enable_tracing` to run again.""" + try: + from ydb.opentelemetry.plugin import _disable_tracing + except ImportError: + return + + _disable_tracing() + + +__all__ = ["disable_tracing", "enable_tracing"] diff --git a/contrib/python/ydb/py3/ydb/opentelemetry/plugin.py b/contrib/python/ydb/py3/ydb/opentelemetry/plugin.py new file mode 100644 index 00000000000..76942789f52 --- /dev/null +++ b/contrib/python/ydb/py3/ydb/opentelemetry/plugin.py @@ -0,0 +1,134 @@ +"""OpenTelemetry bridge for YDB.""" + +from opentelemetry import context as otel_context +from opentelemetry import trace +from opentelemetry.propagate import inject +from opentelemetry.trace import StatusCode + +from ydb import issues +from ydb.issues import StatusCode as YdbStatusCode +from ydb.opentelemetry.tracing import _registry + +# YDB client transport StatusCode values (401xxx band) -> OTel error.type transport_error. +_TRANSPORT_STATUSES = frozenset( + { + YdbStatusCode.CONNECTION_LOST, + YdbStatusCode.CONNECTION_FAILURE, + YdbStatusCode.DEADLINE_EXCEEDED, + YdbStatusCode.CLIENT_INTERNAL_ERROR, + YdbStatusCode.UNIMPLEMENTED, + } +) + +_tracer = None +_enabled = False + +_KIND_MAP = { + "client": trace.SpanKind.CLIENT, + "internal": trace.SpanKind.INTERNAL, +} + + +def _otel_metadata_hook(): + """Inject W3C Trace Context into outgoing gRPC metadata using the active OTel context.""" + headers = {} + inject(headers) + return list(headers.items()) + + +def _set_error_on_span(span, exception): + if isinstance(exception, issues.Error) and exception.status is not None: + span.set_attribute("db.response.status_code", exception.status.name) + error_type = "transport_error" if exception.status in _TRANSPORT_STATUSES else "ydb_error" + else: + error_type = type(exception).__qualname__ + + span.set_attribute("error.type", error_type) + span.set_status(StatusCode.ERROR, str(exception)) + span.record_exception(exception) + + +class _AttachContext: + """Make a span the active OTel context for a ``with`` block. + + When ``end_on_exit=True`` (default) the span is ended on exit — used for + single-shot RPCs. When ``end_on_exit=False`` the span is only ended on + exception — used for streaming RPCs where the result iterator owns ``end()``. + """ + + def __init__(self, span, end_on_exit): + self._span = span + self._end_on_exit = end_on_exit + self._token = None + + def __enter__(self): + ctx = trace.set_span_in_context(self._span._span) + self._token = otel_context.attach(ctx) + return self._span + + def __exit__(self, exc_type, exc_val, exc_tb): + if self._token is not None: + otel_context.detach(self._token) + self._token = None + if exc_val is not None: + self._span.set_error(exc_val) + self._span.end() + elif self._end_on_exit: + self._span.end() + return False + + +class TracingSpan: + """Wrapper around an OTel span. + + Use :meth:`attach_context` as a context manager around any RPC call. + The default (``end_on_exit=True``) is for single-shot operations; pass + ``end_on_exit=False`` for streaming RPCs where the result iterator owns + ``end()``. + """ + + def __init__(self, span): + self._span = span + + def set_error(self, exception): + _set_error_on_span(self._span, exception) + + def set_attribute(self, key, value): + self._span.set_attribute(key, value) + + def end(self): + self._span.end() + + def attach_context(self, end_on_exit=True): + return _AttachContext(self, end_on_exit) + + +def _create_span(name, attributes=None, kind=None): + span = _tracer.start_span( + name, + kind=_KIND_MAP.get(kind, trace.SpanKind.CLIENT), + attributes=attributes or {}, + ) + return TracingSpan(span) + + +def _enable_tracing(tracer=None): + global _enabled, _tracer + + if _enabled: + return + + _tracer = tracer if tracer is not None else trace.get_tracer("ydb.sdk") + _enabled = True + _registry.set_metadata_hook(_otel_metadata_hook) + _registry.set_create_span(_create_span) + + +def _disable_tracing(): + """Clear hooks and tracer; after this, :func:`~ydb.opentelemetry.enable_tracing` may be called again.""" + global _enabled, _tracer + + _registry.set_create_span(None) + _registry.set_metadata_hook(None) + _enabled = False + _tracer = None diff --git a/contrib/python/ydb/py3/ydb/opentelemetry/tracing.py b/contrib/python/ydb/py3/ydb/opentelemetry/tracing.py new file mode 100644 index 00000000000..1d0995df8d5 --- /dev/null +++ b/contrib/python/ydb/py3/ydb/opentelemetry/tracing.py @@ -0,0 +1,165 @@ +"""Internal SDK tracing helpers and registry.""" + +import enum +from typing import Optional, Tuple + + +class SpanName(str, enum.Enum): + """Canonical span names used across the YDB SDK.""" + + CREATE_SESSION = "ydb.CreateSession" + EXECUTE_QUERY = "ydb.ExecuteQuery" + BEGIN_TRANSACTION = "ydb.BeginTransaction" + COMMIT = "ydb.Commit" + ROLLBACK = "ydb.Rollback" + DRIVER_INITIALIZE = "ydb.Driver.Initialize" + RUN_WITH_RETRY = "ydb.RunWithRetry" + TRY = "ydb.Try" + + +class _NoopCtx: + __slots__ = ("_span",) + + def __init__(self, span): + self._span = span + + def __enter__(self): + return self._span + + def __exit__(self, exc_type, exc_val, exc_tb): + return False + + +class _NoopSpan: + """Returned by create_ydb_span when tracing is disabled.""" + + def set_error(self, exception): + pass + + def set_attribute(self, key, value): + pass + + def end(self): + pass + + def attach_context(self, end_on_exit=True): + return _NoopCtx(self) + + +_NOOP_SPAN = _NoopSpan() + + +class OtelTracingRegistry: + """Singleton registry for OpenTelemetry tracing. + + By default everything is no-op until :func:`~ydb.opentelemetry.enable_tracing` is called. + """ + + def __init__(self): + self._metadata_hook = None + self._create_span_func = None + + def is_active(self) -> bool: + return self._create_span_func is not None + + def create_span(self, name, attributes=None, kind=None): + if self._create_span_func is None: + return _NOOP_SPAN + return self._create_span_func(name, attributes, kind=kind) + + def get_trace_metadata(self): + if self._metadata_hook is not None: + return self._metadata_hook() + return [] + + def set_metadata_hook(self, hook): + self._metadata_hook = hook + + def set_create_span(self, func): + self._create_span_func = func + + +_registry = OtelTracingRegistry() + + +def get_trace_metadata(): + """Return tracing metadata for gRPC calls.""" + return _registry.get_trace_metadata() + + +def _split_endpoint(endpoint: Optional[str]) -> Tuple[str, int]: + ep = endpoint or "" + if ep.startswith("grpcs://"): + ep = ep[len("grpcs://") :] + elif ep.startswith("grpc://"): + ep = ep[len("grpc://") :] + + if ep.startswith("["): + close = ep.find("]") + if close != -1 and len(ep) > close + 1 and ep[close + 1] == ":": + host = ep[: close + 1] + port_s = ep[close + 2 :] + return host, int(port_s) if port_s.isdigit() else 0 + + host, sep, port_s = ep.rpartition(":") + if not sep: + return ep, 0 + return host, int(port_s) if port_s.isdigit() else 0 + + +def _build_ydb_attrs(driver_config, node_id=None, peer=None): + host, port = _split_endpoint(getattr(driver_config, "endpoint", None)) + attrs = { + "db.system.name": "ydb", + "db.namespace": getattr(driver_config, "database", None) or "", + "server.address": host, + "server.port": port, + } + if peer is not None: + address, port_, location = peer + if address is not None: + attrs["network.peer.address"] = address + if port_ is not None: + attrs["network.peer.port"] = int(port_) + if location: + attrs["ydb.node.dc"] = location + if node_id is not None: + attrs["ydb.node.id"] = node_id + return attrs + + +def create_span(name, attributes=None, kind="internal"): + """Create a span with no YDB-specific attributes (used for SDK-internal operations).""" + return _registry.create_span(name, attributes=attributes, kind=kind).attach_context() + + +def create_ydb_span(name, driver_config, node_id=None, kind=None, peer=None): + """Create a span pre-filled with standard YDB attributes.""" + if not _registry.is_active(): + return _NOOP_SPAN + attrs = _build_ydb_attrs(driver_config, node_id, peer) + return _registry.create_span(name, attributes=attrs, kind=kind) + + +def set_peer_attributes(span, peer): + """Fill in network.peer.* and ydb.node.dc on an existing span once the peer is known.""" + if peer is None: + return + address, port, location = peer + if address is not None: + span.set_attribute("network.peer.address", address) + if port is not None: + span.set_attribute("network.peer.port", int(port)) + if location: + span.set_attribute("ydb.node.dc", location) + + +def span_finish_callback(span): + """Return an on_finish callable that ends *span* when a streaming result iterator completes.""" + + def _finish(exception=None): + if exception is not None: + span.set_error(exception) + span.end() + + return _finish diff --git a/contrib/python/ydb/py3/ydb/pool.py b/contrib/python/ydb/py3/ydb/pool.py index 2901c573435..f1a4bdf940b 100644 --- a/contrib/python/ydb/py3/ydb/pool.py +++ b/contrib/python/ydb/py3/ydb/pool.py @@ -10,6 +10,7 @@ import random from typing import Any, Callable, ContextManager, List, Optional, Set, Tuple, TYPE_CHECKING from . import connection as connection_impl, issues, resolver, _utilities, tracing +from .opentelemetry.tracing import SpanName, create_ydb_span from abc import abstractmethod from .connection import Connection, EndpointKey @@ -400,23 +401,41 @@ class ConnectionPool(IConnectionPool): self._store = ConnectionsCache(driver_config.use_all_nodes, driver_config.tracer) self.tracer = driver_config.tracer self._grpc_init = connection_impl.Connection(self._driver_config.endpoint, self._driver_config) + self._stopped = False + self._stop_guard = threading.Lock() + self._stop_event = threading.Event() + self._init_thread: Optional[threading.Thread] = None if driver_config.disable_discovery: - # If discovery is disabled, just add the initial endpoint to the store - ready_connection = connection_impl.Connection.ready_factory( - self._driver_config.endpoint, - self._driver_config, - ready_timeout=getattr(self._driver_config, "discovery_request_timeout", 10), - ) - self._store.add(ready_connection) + # If discovery is disabled, establish the initial connection in a + # background thread, retrying until it succeeds or the pool is stopped. + # Doing this off the constructor keeps wait(timeout) as the blocking + # point and lets stop() interrupt the retry loop. self._discovery_thread = None + self._init_thread = threading.Thread( + name="ydb_driver_initial_connection", + target=self._init_connection, + daemon=True, + ) + self._init_thread.start() else: # Start discovery thread as usual self._discovery_thread = Discovery(self._store, self._driver_config) self._discovery_thread.start() - self._stopped = False - self._stop_guard = threading.Lock() + def _init_connection(self) -> None: + ready_timeout = getattr(self._driver_config, "discovery_request_timeout", 10) + while not self._stopped: + ready_connection = connection_impl.Connection.ready_factory( + self._driver_config.endpoint, + self._driver_config, + ready_timeout=ready_timeout, + ) + if self._store.add(ready_connection): + return + + logger.debug("Initial connection attempt failed") + self._stop_event.wait(1) def stop(self, timeout: int = 10) -> None: """ @@ -430,11 +449,16 @@ class ConnectionPool(IConnectionPool): return self._stopped = True + self._stop_event.set() if self._discovery_thread: self._discovery_thread.stop() self._grpc_init.close() if self._discovery_thread: self._discovery_thread.join(timeout) + if self._init_thread: + self._init_thread.join(timeout) + if self._discovery_thread is None: + self._store.cleanup() def async_wait(self, fail_fast: bool = False) -> "futures.Future[None]": """ @@ -453,10 +477,11 @@ class ConnectionPool(IConnectionPool): :param timeout: A timeout to wait in seconds :return: None """ - if fail_fast: - self._store.add_fast_fail().result(timeout) - else: - self._store.subscribe().result(timeout) + with create_ydb_span(SpanName.DRIVER_INITIALIZE, self._driver_config, kind="internal").attach_context(): + if fail_fast: + self._store.add_fast_fail().result(timeout) + else: + self._store.subscribe().result(timeout) def _on_disconnected(self, connection: Connection) -> None: """ diff --git a/contrib/python/ydb/py3/ydb/query/base.py b/contrib/python/ydb/py3/ydb/query/base.py index 7fea6cd0e8c..093c7d55452 100644 --- a/contrib/python/ydb/py3/ydb/query/base.py +++ b/contrib/python/ydb/py3/ydb/query/base.py @@ -27,7 +27,6 @@ from .. import _apis from ydb._topic_common.common import CallFromSyncToAsync, _get_shared_event_loop from ydb._grpc.grpcwrapper.common_utils import to_thread - if typing.TYPE_CHECKING: from .transaction import BaseQueryTxContext from .session import BaseQuerySession @@ -73,9 +72,12 @@ class QueryResultSetFormat(enum.IntEnum): class SyncResponseContextIterator(_utilities.SyncResponseIterator): - def __init__(self, it, wrapper, on_error=None): + """Streams ExecuteQuery results.""" + + def __init__(self, it, wrapper, on_error=None, on_finish=None): super().__init__(it, wrapper) self._on_error = on_error + self._on_finish = on_finish def __enter__(self) -> "SyncResponseContextIterator": return self @@ -86,6 +88,7 @@ class SyncResponseContextIterator(_utilities.SyncResponseIterator): except StopIteration: # Normal stream termination is not an error and must not invalidate # the session. + self._call_on_finish() raise except BaseException as e: # BaseException (not Exception) for parity with the async iterator: @@ -95,8 +98,17 @@ class SyncResponseContextIterator(_utilities.SyncResponseIterator): # SessionBusy. if self._on_error: self._on_error(e) + self._call_on_finish(e) raise + def _call_on_finish(self, exception=None): + if self._on_finish is not None: + self._on_finish(exception) + self._on_finish = None + + def __del__(self): + self._call_on_finish() + def __exit__(self, exc_type, exc_val, exc_tb): # To close stream on YDB it is necessary to scroll through it to the end. # Errors during the cleanup drain have already been reported to _on_error @@ -107,6 +119,7 @@ class SyncResponseContextIterator(_utilities.SyncResponseIterator): pass except BaseException: pass + self._call_on_finish() class QueryClientSettings: diff --git a/contrib/python/ydb/py3/ydb/query/session.py b/contrib/python/ydb/py3/ydb/query/session.py index f2099f8cdbf..a9c1b4a50d9 100644 --- a/contrib/python/ydb/py3/ydb/query/session.py +++ b/contrib/python/ydb/py3/ydb/query/session.py @@ -18,6 +18,7 @@ from . import base from .base import QueryExplainResultFormat from .. import _apis, issues, _utilities +from ..opentelemetry.tracing import SpanName, create_ydb_span, set_peer_attributes, span_finish_callback from ..settings import BaseRequestSettings from ..connection import _RpcState as RpcState, EndpointKey from .._grpc.grpcwrapper import common_utils @@ -30,7 +31,7 @@ from .transaction import QueryTxContext from .._constants import DEFAULT_INITIAL_RESPONSE_TIMEOUT, DEFAULT_LONG_STREAM_TIMEOUT if TYPE_CHECKING: - from ..driver import Driver as SyncDriver + from ..driver import Driver as SyncDriver, DriverConfig from ..aio.driver import Driver as AsyncDriver @@ -46,9 +47,30 @@ def wrapper_create_session( issues._process_response(message.status) session._session_id = message.session_id session._node_id = message.node_id + session._peer = _resolve_peer(session._driver, message.node_id) return session +def _resolve_peer(driver, node_id): + """Look up network.peer.* / ydb.node.dc for a node in the driver's endpoint map.""" + if node_id is None: + return None + store = getattr(driver, "_store", None) + if store is None: + return None + by_node = getattr(store, "connections_by_node_id", None) + if not by_node: + return None + connection = by_node.get(node_id) + if connection is None: + return None + return ( + getattr(connection, "peer_address", None), + getattr(connection, "peer_port", None), + getattr(connection, "peer_location", None), + ) + + def wrapper_delete_session( rpc_state: RpcState, response_pb: _apis.ydb_query.DeleteSessionResponse, @@ -69,6 +91,7 @@ class BaseQuerySession(abc.ABC, Generic[DriverT]): # Session data _session_id: Optional[str] = None _node_id: Optional[int] = None + _peer: Optional[tuple] = None _closed: bool = False _invalidated: bool = False @@ -85,6 +108,10 @@ class BaseQuerySession(abc.ABC, Generic[DriverT]): self._last_query_stats = None @property + def _driver_config(self) -> Optional["DriverConfig"]: + return getattr(self._driver, "_driver_config", None) + + @property def session_id(self) -> Optional[str]: return self._session_id @@ -391,8 +418,10 @@ class QuerySession(BaseQuerySession["SyncDriver"]): if self._closed: raise RuntimeError("Session is already closed.") - self._create_call(settings=settings) - self._attach() + with create_ydb_span(SpanName.CREATE_SESSION, self._driver_config).attach_context() as span: + self._create_call(settings=settings) + set_peer_attributes(span, self._peer) + self._attach() return self @@ -458,20 +487,27 @@ class QuerySession(BaseQuerySession["SyncDriver"]): """ self._check_session_ready_to_use() - stream_it = self._execute_call( - query=query, - parameters=parameters, - commit_tx=True, - syntax=syntax, - exec_mode=exec_mode, - stats_mode=stats_mode, - schema_inclusion_mode=schema_inclusion_mode, - result_set_format=result_set_format, - arrow_format_settings=arrow_format_settings, - concurrent_result_sets=concurrent_result_sets, - settings=settings, + span = create_ydb_span( + SpanName.EXECUTE_QUERY, + self._driver_config, + node_id=self._node_id, + peer=self._peer, ) + with span.attach_context(end_on_exit=False): + stream_it = self._execute_call( + query=query, + parameters=parameters, + commit_tx=True, + syntax=syntax, + exec_mode=exec_mode, + stats_mode=stats_mode, + schema_inclusion_mode=schema_inclusion_mode, + result_set_format=result_set_format, + arrow_format_settings=arrow_format_settings, + concurrent_result_sets=concurrent_result_sets, + settings=settings, + ) return base.SyncResponseContextIterator( stream_it, lambda resp: base.wrap_execute_query_response( @@ -481,6 +517,7 @@ class QuerySession(BaseQuerySession["SyncDriver"]): settings=self._settings, ), on_error=self._on_execute_stream_error, + on_finish=span_finish_callback(span), ) def explain( diff --git a/contrib/python/ydb/py3/ydb/query/transaction.py b/contrib/python/ydb/py3/ydb/query/transaction.py index fdcefb0b801..1d278ac2fee 100644 --- a/contrib/python/ydb/py3/ydb/query/transaction.py +++ b/contrib/python/ydb/py3/ydb/query/transaction.py @@ -17,6 +17,7 @@ from .. import ( _apis, issues, ) +from ..opentelemetry.tracing import SpanName, create_ydb_span, span_finish_callback from .._grpc.grpcwrapper import ydb_topic as _ydb_topic from .._grpc.grpcwrapper import ydb_query as _ydb_query from ..connection import _RpcState as RpcState @@ -245,6 +246,10 @@ class BaseQueryTxContext(base.CallbackHandler, Generic[DriverT]): self._last_query_stats = None @property + def _driver_config(self): + return getattr(self._driver, "_driver_config", None) + + @property def session_id(self) -> Optional[str]: """ A transaction's session id @@ -523,7 +528,13 @@ class QueryTxContext(BaseQueryTxContext["SyncDriver"]): :return: Transaction object or exception if begin is failed """ - self._begin_call(settings) + with create_ydb_span( + SpanName.BEGIN_TRANSACTION, + self._driver_config, + node_id=self.session.node_id, + peer=getattr(self.session, "_peer", None), + ).attach_context(): + self._begin_call(settings) return self @@ -545,13 +556,19 @@ class QueryTxContext(BaseQueryTxContext["SyncDriver"]): self._ensure_prev_stream_finished() - try: - self._execute_callbacks_sync(base.TxEvent.BEFORE_COMMIT) - self._commit_call(settings) - self._execute_callbacks_sync(base.TxEvent.AFTER_COMMIT, exc=None) - except BaseException as e: # TODO: probably should be less wide - self._execute_callbacks_sync(base.TxEvent.AFTER_COMMIT, exc=e) - raise e + with create_ydb_span( + SpanName.COMMIT, + self._driver_config, + node_id=self.session.node_id, + peer=getattr(self.session, "_peer", None), + ).attach_context(): + try: + self._execute_callbacks_sync(base.TxEvent.BEFORE_COMMIT) + self._commit_call(settings) + self._execute_callbacks_sync(base.TxEvent.AFTER_COMMIT, exc=None) + except BaseException as e: # TODO: probably should be less wide + self._execute_callbacks_sync(base.TxEvent.AFTER_COMMIT, exc=e) + raise e def rollback(self, settings: Optional[BaseRequestSettings] = None) -> None: """Calls rollback on a transaction if it is open otherwise is no-op. If transaction execution @@ -571,13 +588,19 @@ class QueryTxContext(BaseQueryTxContext["SyncDriver"]): self._ensure_prev_stream_finished() - try: - self._execute_callbacks_sync(base.TxEvent.BEFORE_ROLLBACK) - self._rollback_call(settings) - self._execute_callbacks_sync(base.TxEvent.AFTER_ROLLBACK, exc=None) - except BaseException as e: # TODO: probably should be less wide - self._execute_callbacks_sync(base.TxEvent.AFTER_ROLLBACK, exc=e) - raise e + with create_ydb_span( + SpanName.ROLLBACK, + self._driver_config, + node_id=self.session.node_id, + peer=getattr(self.session, "_peer", None), + ).attach_context(): + try: + self._execute_callbacks_sync(base.TxEvent.BEFORE_ROLLBACK) + self._rollback_call(settings) + self._execute_callbacks_sync(base.TxEvent.AFTER_ROLLBACK, exc=None) + except BaseException as e: # TODO: probably should be less wide + self._execute_callbacks_sync(base.TxEvent.AFTER_ROLLBACK, exc=e) + raise e def execute( self, @@ -626,20 +649,27 @@ class QueryTxContext(BaseQueryTxContext["SyncDriver"]): """ self._ensure_prev_stream_finished() - stream_it = self._execute_call( - query=query, - commit_tx=commit_tx, - syntax=syntax, - exec_mode=exec_mode, - stats_mode=stats_mode, - schema_inclusion_mode=schema_inclusion_mode, - result_set_format=result_set_format, - arrow_format_settings=arrow_format_settings, - parameters=parameters, - concurrent_result_sets=concurrent_result_sets, - settings=settings, + span = create_ydb_span( + SpanName.EXECUTE_QUERY, + self._driver_config, + node_id=self.session.node_id, + peer=getattr(self.session, "_peer", None), ) + with span.attach_context(end_on_exit=False): + stream_it = self._execute_call( + query=query, + commit_tx=commit_tx, + syntax=syntax, + exec_mode=exec_mode, + stats_mode=stats_mode, + schema_inclusion_mode=schema_inclusion_mode, + result_set_format=result_set_format, + arrow_format_settings=arrow_format_settings, + parameters=parameters, + concurrent_result_sets=concurrent_result_sets, + settings=settings, + ) self._prev_stream = base.SyncResponseContextIterator( stream_it, lambda resp: base.wrap_execute_query_response( @@ -651,5 +681,6 @@ class QueryTxContext(BaseQueryTxContext["SyncDriver"]): settings=self.session._settings, ), on_error=self.session._on_execute_stream_error, + on_finish=span_finish_callback(span), ) return self._prev_stream diff --git a/contrib/python/ydb/py3/ydb/resolver.py b/contrib/python/ydb/py3/ydb/resolver.py index e9e8e6f93f3..1402e542abc 100644 --- a/contrib/python/ydb/py3/ydb/resolver.py +++ b/contrib/python/ydb/py3/ydb/resolver.py @@ -53,7 +53,11 @@ class EndpointInfo(object): ssl_target_name_override = self.address endpoint_options = conn_impl.EndpointOptions( - ssl_target_name_override=ssl_target_name_override, node_id=self.node_id + ssl_target_name_override=ssl_target_name_override, + node_id=self.node_id, + address=self.address, + port=self.port, + location=self.location, ) if self.ipv6_addrs or self.ipv4_addrs: diff --git a/contrib/python/ydb/py3/ydb/retries.py b/contrib/python/ydb/py3/ydb/retries.py index c151e3d2146..4b7c137f39e 100644 --- a/contrib/python/ydb/py3/ydb/retries.py +++ b/contrib/python/ydb/py3/ydb/retries.py @@ -7,6 +7,11 @@ from typing import Any, Callable, Generator, Optional, Union from . import issues from ._errors import check_retriable_error +from .opentelemetry.tracing import SpanName, create_span as _create_span + + +def _try_span_attrs(backoff_ms: Optional[int]): + return {"ydb.retry.backoff_ms": backoff_ms} if backoff_ms is not None else None class BackoffSettings: @@ -129,19 +134,18 @@ def retry_operation_impl( if not retriable_info.is_retriable: raise - skip_yield_error_types = [ + skip_yield_error_types = ( issues.Aborted, issues.BadSession, issues.NotFound, issues.InternalError, - ] + ) - yield_sleep = True - for t in skip_yield_error_types: - if isinstance(e, t): - yield_sleep = False - - if yield_sleep: + if isinstance(e, skip_yield_error_types): + # Skip the inter-attempt sleep but still emit a marker so consumers + # advance per-attempt bookkeeping (e.g. ``ydb.Try`` spans get backoff=0). + yield YdbRetryOperationSleepOpt(0.0) + else: yield YdbRetryOperationSleepOpt(retriable_info.sleep_timeout_seconds) except Exception as e: @@ -159,12 +163,21 @@ def retry_operation_sync( *args: Any, **kwargs: Any, ) -> Any: - opt_generator = retry_operation_impl(callee, retry_settings, *args, **kwargs) - for next_opt in opt_generator: - if isinstance(next_opt, YdbRetryOperationSleepOpt): - time.sleep(next_opt.timeout) - else: - return next_opt.result + backoff_ms: Optional[int] = None + + @functools.wraps(callee) + def traced_callee(*a: Any, **kw: Any) -> Any: + with _create_span(SpanName.TRY, _try_span_attrs(backoff_ms)): + return callee(*a, **kw) + + with _create_span(SpanName.RUN_WITH_RETRY): + for next_opt in retry_operation_impl(traced_callee, retry_settings, *args, **kwargs): + if isinstance(next_opt, YdbRetryOperationSleepOpt): + backoff_ms = int(next_opt.timeout * 1000) + if next_opt.timeout > 0: + time.sleep(next_opt.timeout) + else: + return next_opt.result return None @@ -186,15 +199,20 @@ async def retry_operation_async( # pylint: disable=W1113 Returns awaitable result of coroutine. If retries are not succussful exception is raised. """ - opt_generator = retry_operation_impl(callee, retry_settings, *args, **kwargs) - for next_opt in opt_generator: - if isinstance(next_opt, YdbRetryOperationSleepOpt): - await asyncio.sleep(next_opt.timeout) - else: - try: - return await next_opt.result - except BaseException as e: # pylint: disable=W0703 - next_opt.set_exception(e) + backoff_ms: Optional[int] = None + with _create_span(SpanName.RUN_WITH_RETRY): + for next_opt in retry_operation_impl(callee, retry_settings, *args, **kwargs): + if isinstance(next_opt, YdbRetryOperationSleepOpt): + backoff_ms = int(next_opt.timeout * 1000) + if next_opt.timeout > 0: + await asyncio.sleep(next_opt.timeout) + else: + with _create_span(SpanName.TRY, _try_span_attrs(backoff_ms)) as try_span: + try: + return await next_opt.result + except BaseException as e: # pylint: disable=W0703 + try_span.set_error(e) + next_opt.set_exception(e) return None diff --git a/contrib/python/ydb/py3/ydb/ydb_version.py b/contrib/python/ydb/py3/ydb/ydb_version.py index d7d1d1869be..989938e1be0 100644 --- a/contrib/python/ydb/py3/ydb/ydb_version.py +++ b/contrib/python/ydb/py3/ydb/ydb_version.py @@ -1 +1 @@ -VERSION = "3.28.4" +VERSION = "3.29.6" |