Merge pull request #204 flag for error on truncated response

Author: rekby

tests/aio/test_tx.py

@@ -153,3 +153,48 @@ async def check_transaction(s: ydb.aio.table.Session):
                 assert rs[0].rows[0].cnt == 2
 
         await pool.retry_operation(check_transaction)
+
+
+@pytest.mark.asyncio
+async def test_truncated_response(driver, table_name, table_path):
+    column_types = ydb.BulkUpsertColumns().add_column("id", ydb.PrimitiveType.Int64)
+
+    rows = []
+
+    rows_count = 1100
+    for i in range(rows_count):
+        rows.append({"id": i})
+
+    await driver.table_client.bulk_upsert(table_path, rows, column_types)
+
+    table_client = driver.table_client  # default table client with driver's settings
+    s = table_client.session()
+    await s.create()
+    t = s.transaction()
+
+    result = await t.execute("SELECT * FROM %s" % table_name)
+    assert result[0].truncated
+    assert len(result[0].rows) == 1000
+
+
+@pytest.mark.asyncio
+async def test_truncated_response_deny(driver, table_name, table_path):
+    column_types = ydb.BulkUpsertColumns().add_column("id", ydb.PrimitiveType.Int64)
+
+    rows = []
+
+    rows_count = 1100
+    for i in range(rows_count):
+        rows.append({"id": i})
+
+    await driver.table_client.bulk_upsert(table_path, rows, column_types)
+
+    table_client = ydb.TableClient(
+        driver, ydb.TableClientSettings().with_allow_truncated_result(False)
+    )
+    s = table_client.session()
+    await s.create()
+    t = s.transaction()
+
+    with pytest.raises(ydb.TruncatedResponseError):
+        await t.execute("SELECT * FROM %s" % table_name)

tests/conftest.py

@@ -138,3 +138,8 @@ def create_table(s):
 
         pool.retry_operation_sync(create_table)
     return table_name
+
+
+@pytest.fixture()
+def table_path(database, table_name) -> str:
+    return database + "/" + table_name

tests/table/test_tx.py

@@ -145,3 +145,49 @@ def check_transaction(s: ydb.table.Session):
                 assert rs[0].rows[0].cnt == 2
 
         pool.retry_operation_sync(check_transaction)
+
+
+def test_truncated_response(driver_sync, table_name, table_path):
+    column_types = ydb.BulkUpsertColumns().add_column("id", ydb.PrimitiveType.Int64)
+
+    rows = []
+
+    rows_count = 1100
+    for i in range(rows_count):
+        rows.append({"id": i})
+
+    driver_sync.table_client.bulk_upsert(table_path, rows, column_types)
+
+    table_client = (
+        driver_sync.table_client
+    )  # default table client with driver's settings
+    s = table_client.session()
+    s.create()
+    t = s.transaction()
+
+    result = t.execute("SELECT * FROM %s" % table_name)
+    assert result[0].truncated
+    assert len(result[0].rows) == 1000
+
+
+@pytest.mark.asyncio
+async def test_truncated_response_deny(driver_sync, table_name, table_path):
+    column_types = ydb.BulkUpsertColumns().add_column("id", ydb.PrimitiveType.Int64)
+
+    rows = []
+
+    rows_count = 1100
+    for i in range(rows_count):
+        rows.append({"id": i})
+
+    driver_sync.table_client.bulk_upsert(table_path, rows, column_types)
+
+    table_client = ydb.TableClient(
+        driver_sync, ydb.TableClientSettings().with_allow_truncated_result(False)
+    )
+    s = table_client.session()
+    s.create()
+    t = s.transaction()
+
+    with pytest.raises(ydb.TruncatedResponseError):
+        t.execute("SELECT * FROM %s" % table_name)

ydb/convert.py

@@ -496,5 +496,13 @@ def __init__(self, result_sets_pb, table_client_settings=None):
             _ResultSet.from_message if not make_lazy else _ResultSet.lazy_from_message
         )
         for result_set in result_sets_pb:
-            result_sets.append(initializer(result_set, table_client_settings))
+            result_set = initializer(result_set, table_client_settings)
+            if (
+                result_set.truncated
+                and not table_client_settings._allow_truncated_result
+            ):
+                raise issues.TruncatedResponseError(
+                    "Response for the request was truncated by server"
+                )
+            result_sets.append(result_set)
         super(ResultSets, self).__init__(result_sets)

ydb/issues.py

@@ -52,6 +52,10 @@ def __init__(self, message, issues=None):
         self.message = message
 
 
+class TruncatedResponseError(Error):
+    status = None
+
+
 class ConnectionError(Error):
     status = None
 

ydb/table.py

@@ -1004,6 +1004,7 @@ def __init__(self):
         self._native_json_in_result_sets = False
         self._native_interval_in_result_sets = False
         self._native_timestamp_in_result_sets = False
+        self._allow_truncated_result = True
 
     def with_native_timestamp_in_result_sets(self, enabled):
         # type:(bool) -> ydb.TableClientSettings
@@ -1040,6 +1041,11 @@ def with_lazy_result_sets(self, enabled):
         self._make_result_sets_lazy = enabled
         return self
 
+    def with_allow_truncated_result(self, enabled):
+        # type:(bool) -> ydb.TableClientSettings
+        self._allow_truncated_result = enabled
+        return self
+
 
 class ScanQueryResult(object):
     def __init__(self, result, table_client_settings):