diff --git a/CHANGELOG.md b/CHANGELOG.md index 70bb5bc3..6306aa90 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,6 @@ + +* Query service client support +* Add dunder version to ydb package * OAuth 2.0 token exchange. Allow multiple resource parameters in according to https://www.rfc-editor.org/rfc/rfc8693 ## 3.14.0 ## diff --git a/docker-compose-tls.yml b/docker-compose-tls.yml index 19693705..f0a4b328 100644 --- a/docker-compose-tls.yml +++ b/docker-compose-tls.yml @@ -11,3 +11,4 @@ services: - ./ydb_certs:/ydb_certs environment: - YDB_USE_IN_MEMORY_PDISKS=true + - YDB_ENABLE_COLUMN_TABLES=true diff --git a/docker-compose.yml b/docker-compose.yml index cb37a377..1a466fab 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -8,3 +8,4 @@ services: hostname: localhost environment: - YDB_USE_IN_MEMORY_PDISKS=true + - YDB_ENABLE_COLUMN_TABLES=true diff --git a/examples/query-service/basic_example.py b/examples/query-service/basic_example.py index b355e10c..cfbb3042 100644 --- a/examples/query-service/basic_example.py +++ b/examples/query-service/basic_example.py @@ -82,6 +82,61 @@ def callee(session): pool.retry_operation_sync(callee) + def callee(session: ydb.QuerySessionSync): + query_print = """select $a""" + + print("=" * 50) + print("Check implicit typed parameters") + + values = [ + 1, + 1.0, + True, + "text", + {"4": 8, "15": 16, "23": 42}, + [{"name": "Michael"}, {"surname": "Scott"}], + ] + + for value in values: + print(f"value: {value}") + with session.transaction().execute( + query=query_print, + parameters={"$a": value}, + commit_tx=True, + ) as results: + for result_set in results: + print(f"rows: {str(result_set.rows)}") + + print("=" * 50) + print("Check typed parameters as tuple pair") + + typed_value = ([1, 2, 3], ydb.ListType(ydb.PrimitiveType.Int64)) + print(f"value: {typed_value}") + + with session.transaction().execute( + query=query_print, + parameters={"$a": typed_value}, + commit_tx=True, + ) as results: + for result_set in results: + print(f"rows: {str(result_set.rows)}") + + print("=" * 50) + print("Check typed parameters as ydb.TypedValue") + + typed_value = ydb.TypedValue(111, ydb.PrimitiveType.Int64) + print(f"value: {typed_value}") + + with session.transaction().execute( + query=query_print, + parameters={"$a": typed_value}, + commit_tx=True, + ) as results: + for result_set in results: + print(f"rows: {str(result_set.rows)}") + + pool.retry_operation_sync(callee) + if __name__ == "__main__": main() diff --git a/tests/query/test_query_parameters.py b/tests/query/test_query_parameters.py new file mode 100644 index 00000000..ff033311 --- /dev/null +++ b/tests/query/test_query_parameters.py @@ -0,0 +1,147 @@ +import pytest +import ydb + + +query = """SELECT $a AS value""" + + +def test_select_implicit_int(pool: ydb.QuerySessionPool): + expected_value = 111 + res = pool.execute_with_retries(query, parameters={"$a": expected_value}) + actual_value = res[0].rows[0]["value"] + assert expected_value == actual_value + + +def test_select_implicit_float(pool: ydb.QuerySessionPool): + expected_value = 11.1 + res = pool.execute_with_retries(query, parameters={"$a": expected_value}) + actual_value = res[0].rows[0]["value"] + assert expected_value == pytest.approx(actual_value) + + +def test_select_implicit_bool(pool: ydb.QuerySessionPool): + expected_value = False + res = pool.execute_with_retries(query, parameters={"$a": expected_value}) + actual_value = res[0].rows[0]["value"] + assert expected_value == actual_value + + +def test_select_implicit_str(pool: ydb.QuerySessionPool): + expected_value = "text" + res = pool.execute_with_retries(query, parameters={"$a": expected_value}) + actual_value = res[0].rows[0]["value"] + assert expected_value == actual_value + + +def test_select_implicit_list(pool: ydb.QuerySessionPool): + expected_value = [1, 2, 3] + res = pool.execute_with_retries(query, parameters={"$a": expected_value}) + actual_value = res[0].rows[0]["value"] + assert expected_value == actual_value + + +def test_select_implicit_dict(pool: ydb.QuerySessionPool): + expected_value = {"a": 1, "b": 2} + res = pool.execute_with_retries(query, parameters={"$a": expected_value}) + actual_value = res[0].rows[0]["value"] + assert expected_value == actual_value + + +def test_select_implicit_list_nested(pool: ydb.QuerySessionPool): + expected_value = [{"a": 1}, {"b": 2}] + res = pool.execute_with_retries(query, parameters={"$a": expected_value}) + actual_value = res[0].rows[0]["value"] + assert expected_value == actual_value + + +def test_select_implicit_dict_nested(pool: ydb.QuerySessionPool): + expected_value = {"a": [1, 2, 3], "b": [4, 5]} + res = pool.execute_with_retries(query, parameters={"$a": expected_value}) + actual_value = res[0].rows[0]["value"] + assert expected_value == actual_value + + +def test_select_implicit_custom_type_raises(pool: ydb.QuerySessionPool): + class CustomClass: + pass + + expected_value = CustomClass() + with pytest.raises(ValueError): + pool.execute_with_retries(query, parameters={"$a": expected_value}) + + +def test_select_implicit_empty_list_raises(pool: ydb.QuerySessionPool): + expected_value = [] + with pytest.raises(ValueError): + pool.execute_with_retries(query, parameters={"$a": expected_value}) + + +def test_select_implicit_empty_dict_raises(pool: ydb.QuerySessionPool): + expected_value = {} + with pytest.raises(ValueError): + pool.execute_with_retries(query, parameters={"$a": expected_value}) + + +def test_select_explicit_primitive(pool: ydb.QuerySessionPool): + expected_value = 111 + res = pool.execute_with_retries(query, parameters={"$a": (expected_value, ydb.PrimitiveType.Int64)}) + actual_value = res[0].rows[0]["value"] + assert expected_value == actual_value + + +def test_select_explicit_list(pool: ydb.QuerySessionPool): + expected_value = [1, 2, 3] + type_ = ydb.ListType(ydb.PrimitiveType.Int64) + res = pool.execute_with_retries(query, parameters={"$a": (expected_value, type_)}) + actual_value = res[0].rows[0]["value"] + assert expected_value == actual_value + + +def test_select_explicit_dict(pool: ydb.QuerySessionPool): + expected_value = {"key": "value"} + type_ = ydb.DictType(ydb.PrimitiveType.Utf8, ydb.PrimitiveType.Utf8) + res = pool.execute_with_retries(query, parameters={"$a": (expected_value, type_)}) + actual_value = res[0].rows[0]["value"] + assert expected_value == actual_value + + +def test_select_explicit_empty_list_not_raises(pool: ydb.QuerySessionPool): + expected_value = [] + type_ = ydb.ListType(ydb.PrimitiveType.Int64) + res = pool.execute_with_retries(query, parameters={"$a": (expected_value, type_)}) + actual_value = res[0].rows[0]["value"] + assert expected_value == actual_value + + +def test_select_explicit_empty_dict_not_raises(pool: ydb.QuerySessionPool): + expected_value = {} + type_ = ydb.DictType(ydb.PrimitiveType.Utf8, ydb.PrimitiveType.Utf8) + res = pool.execute_with_retries(query, parameters={"$a": (expected_value, type_)}) + actual_value = res[0].rows[0]["value"] + assert expected_value == actual_value + + +def test_select_typedvalue_full_primitive(pool: ydb.QuerySessionPool): + expected_value = 111 + typed_value = ydb.TypedValue(expected_value, ydb.PrimitiveType.Int64) + res = pool.execute_with_retries(query, parameters={"$a": typed_value}) + actual_value = res[0].rows[0]["value"] + assert expected_value == actual_value + + +def test_select_typedvalue_implicit_primitive(pool: ydb.QuerySessionPool): + expected_value = 111 + typed_value = ydb.TypedValue(expected_value) + res = pool.execute_with_retries(query, parameters={"$a": typed_value}) + actual_value = res[0].rows[0]["value"] + assert expected_value == actual_value + + +def test_select_typevalue_custom_type_raises(pool: ydb.QuerySessionPool): + class CustomClass: + pass + + expected_value = CustomClass() + typed_value = ydb.TypedValue(expected_value) + with pytest.raises(ValueError): + pool.execute_with_retries(query, parameters={"$a": typed_value}) diff --git a/ydb/_grpc/grpcwrapper/ydb_query.py b/ydb/_grpc/grpcwrapper/ydb_query.py index befb02c7..913b8480 100644 --- a/ydb/_grpc/grpcwrapper/ydb_query.py +++ b/ydb/_grpc/grpcwrapper/ydb_query.py @@ -18,6 +18,8 @@ ServerStatus, ) +from ... import convert + @dataclass class CreateSessionResponse(IFromProto): @@ -176,5 +178,5 @@ def to_proto(self) -> ydb_query_pb2.ExecuteQueryRequest: exec_mode=self.exec_mode, stats_mode=self.stats_mode, concurrent_result_sets=self.concurrent_result_sets, - parameters=self.parameters, + parameters=convert.query_parameters_to_pb(self.parameters), ) diff --git a/ydb/convert.py b/ydb/convert.py index 6c4164bc..63a5dbe4 100644 --- a/ydb/convert.py +++ b/ydb/convert.py @@ -281,6 +281,66 @@ def parameters_to_pb(parameters_types, parameters_values): return param_values_pb +def query_parameters_to_pb(parameters): + if parameters is None or not parameters: + return {} + + parameters_types = {} + parameters_values = {} + for name, value in parameters.items(): + if isinstance(value, types.TypedValue): + if value.value_type is None: + value.value_type = _type_from_python_native(value.value) + elif isinstance(value, tuple): + value = types.TypedValue(*value) + else: + value = types.TypedValue(value, _type_from_python_native(value)) + + parameters_values[name] = value.value + parameters_types[name] = value.value_type + + return parameters_to_pb(parameters_types, parameters_values) + + +_from_python_type_map = { + int: types.PrimitiveType.Int64, + float: types.PrimitiveType.Float, + bool: types.PrimitiveType.Bool, + str: types.PrimitiveType.Utf8, +} + + +def _type_from_python_native(value): + t = type(value) + + if t in _from_python_type_map: + return _from_python_type_map[t] + + if t == list: + if len(value) == 0: + raise ValueError( + "Could not map empty list to any type, please specify " + "it manually by tuple(value, type) or ydb.TypedValue" + ) + entry_type = _type_from_python_native(value[0]) + return types.ListType(entry_type) + + if t == dict: + if len(value) == 0: + raise ValueError( + "Could not map empty dict to any type, please specify " + "it manually by tuple(value, type) or ydb.TypedValue" + ) + entry = list(value.items())[0] + key_type = _type_from_python_native(entry[0]) + value_type = _type_from_python_native(entry[1]) + return types.DictType(key_type, value_type) + + raise ValueError( + "Could not map value to any type, please specify it manually by tuple(value, type) or ydb.TypedValue" + ) + + def _unwrap_optionality(column): c_type = column.type current_type = c_type.WhichOneof("type") diff --git a/ydb/types.py b/ydb/types.py index 2a2a7e07..f8a56e4d 100644 --- a/ydb/types.py +++ b/ydb/types.py @@ -2,6 +2,7 @@ from __future__ import annotations import abc +from dataclasses import dataclass import enum import json from . import _utilities, _apis @@ -441,3 +442,9 @@ def proto(self): def __str__(self): return "BulkUpsertColumns<%s>" % ",".join(self.__columns_repr) + + +@dataclass +class TypedValue: + value: typing.Any + value_type: typing.Optional[typing.Union[PrimitiveType, AbstractTypeBuilder]] = None