diff --git a/tests/test_connect_options.py b/tests/test_connect_options.py new file mode 100644 index 0000000..e666523 --- /dev/null +++ b/tests/test_connect_options.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +import pytest +import ydb_dbapi as dbapi +from ydb_dbapi.errors import ProgrammingError +from ydb_dbapi.utils import prepare_driver_config_kwargs + + +class TestPrepareDriverConfigKwargs: + """Leftover connect() keywords are routed, not dropped.""" + + def test_known_driver_option_is_routed(self): + assert prepare_driver_config_kwargs( + None, {"disable_discovery": True} + ) == {"disable_discovery": True} + + def test_explicit_driver_config_kwargs_are_kept(self): + result = prepare_driver_config_kwargs( + {"grpc_keep_alive_timeout": 777}, {"disable_discovery": True} + ) + assert result == { + "grpc_keep_alive_timeout": 777, + "disable_discovery": True, + } + + def test_no_leftovers_returns_driver_config_kwargs(self): + assert prepare_driver_config_kwargs({"use_all_nodes": False}, {}) == { + "use_all_nodes": False + } + + def test_unknown_option_raises(self): + with pytest.raises(ProgrammingError, match="disable_discovry"): + prepare_driver_config_kwargs(None, {"disable_discovry": True}) + + def test_unknown_option_lists_supported_ones(self): + with pytest.raises(ProgrammingError, match="disable_discovery"): + prepare_driver_config_kwargs(None, {"nonsense": 1}) + + def test_all_unknown_options_are_reported(self): + with pytest.raises(ProgrammingError, match="first, second"): + prepare_driver_config_kwargs(None, {"second": 1, "first": 2}) + + def test_option_reserved_by_the_connection_is_unknown(self): + # The connection computes these itself, so they must not be + # overridable through connect() keywords. + with pytest.raises(ProgrammingError, match="endpoint"): + prepare_driver_config_kwargs(None, {"endpoint": "grpc://host:1"}) + + def test_option_passed_twice_raises(self): + with pytest.raises(ProgrammingError, match="disable_discovery"): + prepare_driver_config_kwargs( + {"disable_discovery": True}, {"disable_discovery": False} + ) + + +class TestDriverOptionCoercion: + """URL query parameters arrive as strings and must not stay strings.""" + + @pytest.mark.parametrize("value", ["true", "True", "yes", "on", "1"]) + def test_truthy_strings(self, value: str): + result = prepare_driver_config_kwargs( + None, {"disable_discovery": value} + ) + assert result["disable_discovery"] is True + + @pytest.mark.parametrize("value", ["false", "False", "no", "off", "0"]) + def test_falsy_strings(self, value: str): + # The whole point of coercion: "false" is a non-empty string and + # would otherwise disable discovery. + result = prepare_driver_config_kwargs( + None, {"disable_discovery": value} + ) + assert result["disable_discovery"] is False + + def test_invalid_boolean_raises(self): + with pytest.raises(ProgrammingError, match="expects a boolean"): + prepare_driver_config_kwargs(None, {"disable_discovery": "maybe"}) + + def test_integer_string(self): + result = prepare_driver_config_kwargs( + None, {"discovery_request_timeout": "42"} + ) + assert result["discovery_request_timeout"] == 42 + + def test_optional_integer_string(self): + result = prepare_driver_config_kwargs( + None, {"grpc_keep_alive_timeout": "777"} + ) + assert result["grpc_keep_alive_timeout"] == 777 + + def test_invalid_integer_raises(self): + with pytest.raises(ProgrammingError, match="expects an integer"): + prepare_driver_config_kwargs( + None, {"discovery_request_timeout": "soon"} + ) + + def test_string_option_is_left_alone(self): + result = prepare_driver_config_kwargs( + None, {"grpc_lb_policy_name": "pick_first"} + ) + assert result["grpc_lb_policy_name"] == "pick_first" + + def test_non_string_values_are_passed_through(self): + result = prepare_driver_config_kwargs( + None, {"discovery_request_timeout": 42} + ) + assert result["discovery_request_timeout"] == 42 + + def test_option_not_expressible_as_string_raises(self): + with pytest.raises(ProgrammingError, match="driver_config_kwargs"): + prepare_driver_config_kwargs(None, {"channel_options": "a=b"}) + + +class TestConnectRejectsUnusableOptions: + """connect() reports bad options before touching the network.""" + + def test_unknown_keyword(self, connection_kwargs: dict): + with pytest.raises(ProgrammingError, match="disable_discovry"): + dbapi.connect(**connection_kwargs, disable_discovry=True) + + def test_auth_token_together_with_credentials( + self, connection_kwargs: dict + ): + # auth_token silently wins inside DriverConfig, so refuse the + # combination instead of dropping the credentials. + with pytest.raises(ProgrammingError, match="auth_token"): + dbapi.connect( + **connection_kwargs, + credentials={"token": "some-token"}, + auth_token="another-token", + ) + + def test_driver_option_with_shared_session_pool( + self, connection_kwargs: dict + ): + with pytest.raises(ProgrammingError, match="ydb_session_pool"): + dbapi.connect( + **connection_kwargs, + ydb_session_pool=object(), + disable_discovery=True, + ) + + def test_driver_config_kwargs_with_shared_session_pool( + self, connection_kwargs: dict + ): + with pytest.raises(ProgrammingError, match="ydb_session_pool"): + dbapi.connect( + **connection_kwargs, + ydb_session_pool=object(), + driver_config_kwargs={"disable_discovery": True}, + ) + + def test_additional_sdk_headers_are_still_accepted( + self, connection_kwargs: dict + ): + # Private channel used by ydb-sqlalchemy: it must never be + # treated as a driver option. + conn = dbapi.Connection( + **connection_kwargs, + _additional_sdk_headers=("ydb-sqlalchemy/0.0.0",), + ) + try: + headers = conn._driver._driver_config._additional_sdk_headers + assert "ydb-sqlalchemy/0.0.0" in headers + finally: + conn.close() + + def test_additional_sdk_headers_with_shared_session_pool( + self, connection_kwargs: dict + ): + # It is not a driver option, so it must not trip the shared pool + # check either. + with pytest.raises(AttributeError): + dbapi.Connection( + **connection_kwargs, + ydb_session_pool=object(), + _additional_sdk_headers=("ydb-sqlalchemy/0.0.0",), + ) diff --git a/tests/test_connections.py b/tests/test_connections.py index 029a064..75e9bec 100644 --- a/tests/test_connections.py +++ b/tests/test_connections.py @@ -460,6 +460,25 @@ def test_connect_with_custom_driver_config_kwargs( assert connection._driver._driver_config.grpc_keep_alive_timeout == 777 connection.close() + def test_connect_with_driver_option_as_keyword( + self, connection_kwargs: dict + ) -> None: + connection = dbapi.connect( + **connection_kwargs, grpc_keep_alive_timeout=777 + ) + assert connection._driver._driver_config.grpc_keep_alive_timeout == 777 + connection.close() + + def test_connect_with_driver_option_as_string( + self, connection_kwargs: dict + ) -> None: + # How a SQLAlchemy URL query parameter arrives. + connection = dbapi.connect( + **connection_kwargs, disable_discovery="false" + ) + assert connection._driver._driver_config.disable_discovery is False + connection.close() + @pytest.mark.parametrize( ("isolation_level", "read_only"), [ diff --git a/ydb_dbapi/connections.py b/ydb_dbapi/connections.py index 285fd1b..672da91 100644 --- a/ydb_dbapi/connections.py +++ b/ydb_dbapi/connections.py @@ -19,9 +19,11 @@ from .errors import InterfaceError from .errors import InternalError from .errors import NotSupportedError +from .errors import ProgrammingError from .utils import handle_ydb_errors from .utils import maybe_get_current_trace_id from .utils import prepare_credentials +from .utils import prepare_driver_config_kwargs from .version import VERSION @@ -94,7 +96,25 @@ def __init__( self.connection_kwargs: dict = kwargs - driver_config_kwargs = driver_config_kwargs or {} + # Reserved for SDK integrations rather than a user-facing option, + # so it is taken out untouched before the remaining keywords are + # validated as driver ones. + _additional_sdk_headers: tuple[str, ...] = () + if "_additional_sdk_headers" in kwargs: + val = kwargs.pop("_additional_sdk_headers") + if isinstance(val, tuple): + _additional_sdk_headers = val + + if "auth_token" in kwargs and self.credentials is not None: + msg = ( + "Both credentials and auth_token are set: auth_token " + "would silently override credentials." + ) + raise ProgrammingError(msg) + + driver_config_kwargs = prepare_driver_config_kwargs( + driver_config_kwargs, kwargs + ) self._shared_session_pool: bool = False @@ -103,6 +123,15 @@ def __init__( self.interactive_transaction: bool = False if ydb_session_pool is not None: + if driver_config_kwargs: + names = ", ".join(sorted(driver_config_kwargs)) + msg = ( + f"Driver option(s) {names} cannot be applied: " + "ydb_session_pool comes with its own driver. " + "Configure the driver before creating the pool." + ) + raise ProgrammingError(msg) + self._shared_session_pool = True self._session_pool = ydb_session_pool settings = self._get_client_settings() @@ -114,12 +143,6 @@ def __init__( root_certificates_path ) - _additional_sdk_headers: tuple[str, ...] = () - if "_additional_sdk_headers" in kwargs: - val = kwargs.pop("_additional_sdk_headers") - if isinstance(val, tuple): - _additional_sdk_headers = val - framework_headers = ( f"ydb-dbapi/{VERSION}", *_additional_sdk_headers, diff --git a/ydb_dbapi/utils.py b/ydb_dbapi/utils.py index 3f7d212..1c456e7 100644 --- a/ydb_dbapi/utils.py +++ b/ydb_dbapi/utils.py @@ -7,9 +7,14 @@ import json import re from enum import Enum +from inspect import Parameter from inspect import iscoroutinefunction +from inspect import signature from typing import Any from typing import Callable +from typing import Union +from typing import get_args +from typing import get_origin import ydb @@ -169,6 +174,128 @@ def prepare_credentials( return ydb.AnonymousCredentials() +# Built by the connection itself: accepting them from the caller would +# silently override the value the connection is constructed with. +_RESERVED_DRIVER_CONFIG_PARAMS = frozenset( + { + "endpoint", + "database", + "credentials", + "root_certificates", + "query_client_settings", + "_additional_sdk_headers", + } +) + +_TRUE_STRINGS = frozenset({"1", "on", "true", "yes"}) +_FALSE_STRINGS = frozenset({"0", "off", "false", "no"}) + + +def _driver_config_params() -> dict[str, Parameter]: + # Ask the SDK instead of keeping a list here: DriverConfig gains + # parameters over time, and a stale copy would reject valid ones. + params = signature(ydb.DriverConfig.__init__).parameters + return { + name: param + for name, param in params.items() + if name != "self" and name not in _RESERVED_DRIVER_CONFIG_PARAMS + } + + +def _driver_config_value_type(param: Parameter) -> type | None: + annotation = param.annotation + if get_origin(annotation) is Union: + args = [a for a in get_args(annotation) if a is not type(None)] + if len(args) == 1: + annotation = args[0] + + if annotation in (bool, int, str): + return annotation + + # Fallback for absent or stringified annotations. bool goes first, + # since bool is a subclass of int. + for candidate in (bool, int, str): + if isinstance(param.default, candidate): + return candidate + + return None + + +def _coerce_driver_config_value( + name: str, value: Any, param: Parameter +) -> Any: + # Only strings are coerced. Query parameters of a SQLAlchemy URL + # always arrive as strings, so `disable_discovery=false` would + # otherwise be a non-empty, and therefore truthy, string. + if not isinstance(value, str): + return value + + value_type = _driver_config_value_type(param) + + if value_type is bool: + if value.lower() in _TRUE_STRINGS: + return True + if value.lower() in _FALSE_STRINGS: + return False + msg = f"Connection option {name!r} expects a boolean, got {value!r}." + raise ProgrammingError(msg) + + if value_type is int: + try: + return int(value) + except ValueError: + msg = ( + f"Connection option {name!r} expects an integer, " + f"got {value!r}." + ) + raise ProgrammingError(msg) from None + + if value_type is str: + return value + + msg = ( + f"Connection option {name!r} cannot be set from a string. " + "Pass it via the driver_config_kwargs argument of connect()." + ) + raise ProgrammingError(msg) + + +def prepare_driver_config_kwargs( + driver_config_kwargs: dict[str, Any] | None, + connect_kwargs: dict[str, Any], +) -> dict[str, Any]: + """Route leftover connect() keywords into ydb.DriverConfig kwargs. + + Whatever the driver does not accept is reported rather than dropped: + an option that never arrived must not be indistinguishable from one + that did not help. + """ + known = _driver_config_params() + + unknown = sorted(name for name in connect_kwargs if name not in known) + if unknown: + msg = ( + f"Unknown connection option(s): {', '.join(unknown)}. " + f"Supported driver options: {', '.join(sorted(known))}." + ) + raise ProgrammingError(msg) + + result = dict(driver_config_kwargs or {}) + + duplicated = sorted(name for name in connect_kwargs if name in result) + if duplicated: + msg = ( + f"Connection option(s) {', '.join(duplicated)} passed both " + "directly and via driver_config_kwargs." + ) + raise ProgrammingError(msg) + + for name, value in connect_kwargs.items(): + result[name] = _coerce_driver_config_value(name, value, known[name]) + + return result + + # Order matters: bool before int, datetime before date (subclass checks). _PYTHON_TO_YDB_TYPE: list[tuple[type, Any]] = [ (bool, ydb.PrimitiveType.Bool),