Ability to use implicit lists and dicts

Author: vgvoleg

examples/query-service/basic_example.py

@@ -16,15 +16,15 @@ def main():
 
     pool = ydb.QuerySessionPool(driver)
 
-    # print("=" * 50)
-    # print("DELETE TABLE IF EXISTS")
-    # pool.execute_with_retries("drop table if exists example")
+    print("=" * 50)
+    print("DELETE TABLE IF EXISTS")
+    pool.execute_with_retries("drop table if exists example")
 
-    # print("=" * 50)
-    # print("CREATE TABLE")
-    # pool.execute_with_retries("CREATE TABLE example(key UInt64, value String, PRIMARY KEY (key))")
+    print("=" * 50)
+    print("CREATE TABLE")
+    pool.execute_with_retries("CREATE TABLE example(key UInt64, value String, PRIMARY KEY (key))")
 
-    # pool.execute_with_retries("INSERT INTO example (key, value) VALUES (1, 'onepieceisreal')")
+    pool.execute_with_retries("INSERT INTO example (key, value) VALUES (1, 'onepieceisreal')")
 
     def callee(session):
         print("=" * 50)
@@ -94,7 +94,7 @@ def callee(session: ydb.QuerySessionSync):
 
         for value in values:
             print(f"value: {value}")
-            with session.transaction().execute(query=query_print, parameters={'$a': value}, commit_tx=True) as results:
+            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)}")
 

tests/query/test_query_parameters.py

@@ -7,50 +7,78 @@
 
 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']
+    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']
+    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']
+    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']
+    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_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']
+    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']
+    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'}
+    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']
+    res = pool.execute_with_retries(query, parameters={"$a": (expected_value, type_)})
+    actual_value = res[0].rows[0]["value"]
     assert expected_value == actual_value

ydb/convert.py

@@ -301,17 +301,33 @@ def query_parameters_to_pb(parameters):
 _from_python_type_map = {
     int: types.PrimitiveType.Int64,
     float: types.PrimitiveType.Float,
-    dict: types.PrimitiveType.Json,
     bool: types.PrimitiveType.Bool,
     str: types.PrimitiveType.Utf8,
 }
 
 
 def _primitive_type_from_python_native(value):
     t = type(value)
-    if t not in _from_python_type_map:
-        return types.PrimitiveType.Int64
-    return _from_python_type_map[t]
+    default_type = types.PrimitiveType.Int64
+
+    if t in _from_python_type_map:
+        return _from_python_type_map[t]
+
+    if t == list:
+        if len(value) == 0:
+            return types.ListType(default_type)
+        entry_type = _primitive_type_from_python_native(value[0])
+        return types.ListType(entry_type)
+
+    if t == dict:
+        if len(value) == 0:
+            return types.DictType(default_type, default_type)
+        entry = list(value.items())[0]
+        key_type = _primitive_type_from_python_native(entry[0])
+        value_type = _primitive_type_from_python_native(entry[1])
+        return types.DictType(key_type, value_type)
+
+    return default_type
 
 
 def _unwrap_optionality(column):