diff --git a/src/Ydb.Sdk/src/Services/Query/QueryClient.cs b/src/Ydb.Sdk/src/Services/Query/QueryClient.cs index 1aaeca6b..38533cb9 100644 --- a/src/Ydb.Sdk/src/Services/Query/QueryClient.cs +++ b/src/Ydb.Sdk/src/Services/Query/QueryClient.cs @@ -231,12 +231,23 @@ internal static ExecuteQueryResponsePart FromProto(Ydb.Query.ExecuteQueryRespons } public class ExecuteQueryStream : StreamResponse - , IAsyncEnumerable, IAsyncEnumerator + , IAsyncEnumerable { internal ExecuteQueryStream(Driver.StreamIterator iterator) : base(iterator) { } + public new async Task Next() + { + var isNext = await base.Next(); + if (isNext) + { + Response.EnsureSuccess(); + } + + return isNext; + } + protected override ExecuteQueryResponsePart MakeResponse(Ydb.Query.ExecuteQueryResponsePart protoResponse) { return ExecuteQueryResponsePart.FromProto(protoResponse); @@ -247,23 +258,14 @@ protected override ExecuteQueryResponsePart MakeResponse(Status status) return new ExecuteQueryResponsePart(status); } - public IAsyncEnumerator GetAsyncEnumerator( + public async IAsyncEnumerator GetAsyncEnumerator( CancellationToken cancellationToken = new CancellationToken()) { - throw new NotImplementedException(); - } - - public ValueTask DisposeAsync() - { - throw new NotImplementedException(); - } - - public ValueTask MoveNextAsync() - { - throw new NotImplementedException(); + while (await Next()) + { + yield return Response; + } } - - public ExecuteQueryResponsePart Current { get; } } public class BeginTransactionResponse : ResponseBase @@ -410,7 +412,7 @@ public async Task DeleteSession(string sessionId, DeleteS public SessionStateStream AttachSession(string sessionId, AttachSessionSettings? settings = null) { - settings ??= new AttachSessionSettings(); + settings ??= new AttachSessionSettings { TransportTimeout = TimeSpan.FromDays(1) }; var request = new AttachSessionRequest { SessionId = sessionId }; diff --git a/src/Ydb.Sdk/tests/Query/TestExecuteQuery.cs b/src/Ydb.Sdk/tests/Query/TestExecuteQuery.cs index 9fd1d29d..051cdcd8 100644 --- a/src/Ydb.Sdk/tests/Query/TestExecuteQuery.cs +++ b/src/Ydb.Sdk/tests/Query/TestExecuteQuery.cs @@ -107,10 +107,8 @@ public async Task Tx() func: async stream => { var titles = new List(); - while (await stream.Next()) + await foreach (var part in stream) { - var part = stream.Response; - part.EnsureSuccess(); var resultSet = part.ResultSet; if (resultSet is not null) { @@ -332,7 +330,6 @@ REPLACE INTO episodes { while (await stream.Next()) { - stream.Response.EnsureSuccess(); } } ); @@ -342,11 +339,8 @@ REPLACE INTO episodes private static async Task> ReadSeries(ExecuteQueryStream stream) { var series = new List(); - // await foreach (var part in stream) // TODO - while (await stream.Next()) + await foreach (var part in stream) { - var part = stream.Response; - part.EnsureSuccess(); var resultSet = part.ResultSet; if (resultSet is not null) { @@ -354,8 +348,6 @@ private static async Task> ReadSeries(ExecuteQueryStream stream) } } - stream.Response.EnsureSuccess(); - return series; } } \ No newline at end of file