diff --git a/CHANGELOG.md b/CHANGELOG.md index c5de91ed..cdf43b67 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,5 @@ * Fixed `OSError: [Errno 22] Invalid argument` when reading native `Datetime64` value before 1970 (negative Unix timestamps) on Windows: conversion now uses epoch arithmetic instead of `datetime.utcfromtimestamp` +* Add `return_not_null_data_as_optional` parameter to `read_table` (sync and async) — allows reading `NOT NULL` columns as non-optional types, matching the existing gRPC field and the C++ SDK's `ReturnNotNullAsOptional`. By default `read_table` still wraps `NOT NULL` columns in `Optional` for backward compatibility ## 3.31.4 ## * Fixed async `QuerySessionPool` permanently losing a pool slot when `acquire()` was cancelled while a new session was being created: `asyncio.CancelledError` no longer leaks the pool size counter, so a pool under deadline-driven cancellations can no longer end up exhausted and blocking forever. A cancelled or interrupted session attach now also closes the session instead of orphaning it server-side diff --git a/ydb/_session_impl.py b/ydb/_session_impl.py index 9d5f7770..8d9ef076 100644 --- a/ydb/_session_impl.py +++ b/ydb/_session_impl.py @@ -353,6 +353,7 @@ def read_table_request_factory( ordered=False, row_limit=None, use_snapshot=None, + return_not_null_data_as_optional=None, ): request = _apis.ydb_table.ReadTableRequest() request.path = path @@ -383,6 +384,14 @@ def read_table_request_factory( request.use_snapshot = _apis.FeatureFlag.DISABLED else: request.use_snapshot = use_snapshot + if return_not_null_data_as_optional is not None: + if isinstance(return_not_null_data_as_optional, bool): + if return_not_null_data_as_optional: + request.return_not_null_data_as_optional = _apis.FeatureFlag.ENABLED + else: + request.return_not_null_data_as_optional = _apis.FeatureFlag.DISABLED + else: + request.return_not_null_data_as_optional = return_not_null_data_as_optional return session_state.attach_request(request) diff --git a/ydb/aio/table.py b/ydb/aio/table.py index f76a934b..70db2d87 100644 --- a/ydb/aio/table.py +++ b/ydb/aio/table.py @@ -44,6 +44,7 @@ async def read_table( row_limit=None, settings=None, use_snapshot=None, + return_not_null_data_as_optional=None, ): # pylint: disable=W0236 request = _session_impl.read_table_request_factory( self._state, @@ -53,6 +54,7 @@ async def read_table( ordered, row_limit, use_snapshot=use_snapshot, + return_not_null_data_as_optional=return_not_null_data_as_optional, ) stream_it = await self._driver( request, diff --git a/ydb/table.py b/ydb/table.py index 8fbc5e78..348caeb0 100644 --- a/ydb/table.py +++ b/ydb/table.py @@ -1053,6 +1053,7 @@ def read_table( row_limit=None, settings=None, use_snapshot=None, + return_not_null_data_as_optional=None, ): """ Perform an read table request. @@ -1772,6 +1773,7 @@ def read_table( row_limit=None, settings=None, use_snapshot=None, + return_not_null_data_as_optional=None, ): """ Perform an read table request. @@ -1794,6 +1796,7 @@ def read_table( ordered, row_limit, use_snapshot=use_snapshot, + return_not_null_data_as_optional=return_not_null_data_as_optional, ) stream_it = self._driver( request, @@ -2032,6 +2035,7 @@ def async_read_table( row_limit=None, settings=None, use_snapshot=None, + return_not_null_data_as_optional=None, ): """ Perform an read table request. @@ -2056,6 +2060,7 @@ def async_read_table( ordered, row_limit, use_snapshot=use_snapshot, + return_not_null_data_as_optional=return_not_null_data_as_optional, ) stream_it = self._driver( request, diff --git a/ydb/table_test.py b/ydb/table_test.py index ad158260..23a71413 100644 --- a/ydb/table_test.py +++ b/ydb/table_test.py @@ -4,7 +4,7 @@ from unittest import mock from . import issues, convert, types, _apis, scheme, _session_impl -from .table import SystemViewSchemeEntry, TableClient +from .table import SystemViewSchemeEntry, TableClient, TableClientSettings from .retries import ( retry_operation_impl, @@ -336,3 +336,38 @@ def future(self, request, stub, method, wrap_fn, settings, wrap_args, *rest): assert driver.request.path == "/local/.sys/partition_stats" assert isinstance(entry, SystemViewSchemeEntry) assert entry.sys_view_name == "partition_stats" + + +def _read_table_session_state(): + state = _session_impl.SessionState(TableClientSettings()) + state.set_id("test-session-id") + return state + + +def test_read_table_request_not_null_as_optional_enabled(): + request = _session_impl.read_table_request_factory( + _read_table_session_state(), "/local/table", return_not_null_data_as_optional=True + ) + assert request.return_not_null_data_as_optional == _apis.FeatureFlag.ENABLED + + +def test_read_table_request_not_null_as_optional_disabled(): + request = _session_impl.read_table_request_factory( + _read_table_session_state(), "/local/table", return_not_null_data_as_optional=False + ) + assert request.return_not_null_data_as_optional == _apis.FeatureFlag.DISABLED + + +def test_read_table_request_not_null_as_optional_unset_by_default(): + request = _session_impl.read_table_request_factory(_read_table_session_state(), "/local/table") + assert request.return_not_null_data_as_optional == _apis.FeatureFlag.STATUS_UNSPECIFIED + + +def test_read_table_request_not_null_as_optional_raw_status(): + # A raw FeatureFlag.Status (not a bool) is passed through unchanged. + request = _session_impl.read_table_request_factory( + _read_table_session_state(), + "/local/table", + return_not_null_data_as_optional=_apis.FeatureFlag.DISABLED, + ) + assert request.return_not_null_data_as_optional == _apis.FeatureFlag.DISABLED