Merge pull request #463 from ydb-platform/typed_parameters

vgvoleg · web-flow · commit 13c5985bf03d · 2024-08-02T18:37:03.000+03:00
Typed parameters in Query Service
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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

Original file line number	Diff line number	Diff line change
`@@ -18,6 +18,8 @@`
`18`	`18`	`ServerStatus,`
`19`	`19`	`)`
`20`	`20`
	`21`	`+from ... import convert`
	`22`	`+`
`21`	`23`
`22`	`24`	`@dataclass`
`23`	`25`	`class CreateSessionResponse(IFromProto):`
`@@ -176,5 +178,5 @@ def to_proto(self) -> ydb_query_pb2.ExecuteQueryRequest:`
`176`	`178`	`exec_mode=self.exec_mode,`
`177`	`179`	`stats_mode=self.stats_mode,`
`178`	`180`	`concurrent_result_sets=self.concurrent_result_sets,`
`179`		`- parameters=self.parameters,`
	`181`	`+ parameters=convert.query_parameters_to_pb(self.parameters),`
`180`	`182`	`)`