make ExecuteQueryStream implement IAsyncEnumerable

ydb-platform · Nov 1, 2023 · 221b25a · 221b25a
1 parent 75df84b
commit 221b25a
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 26 deletions.
diff --git 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<Ydb.Query.ExecuteQueryResponsePart, ExecuteQueryResponsePart>
-    , IAsyncEnumerable<ExecuteQueryResponsePart>, IAsyncEnumerator<ExecuteQueryResponsePart>
+    , IAsyncEnumerable<ExecuteQueryResponsePart>
 {
     internal ExecuteQueryStream(Driver.StreamIterator<Ydb.Query.ExecuteQueryResponsePart> iterator) : base(iterator)
     {
     }
 
+    public new async Task<bool> 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<ExecuteQueryResponsePart> GetAsyncEnumerator(
+    public async IAsyncEnumerator<ExecuteQueryResponsePart> GetAsyncEnumerator(
         CancellationToken cancellationToken = new CancellationToken())
     {
-        throw new NotImplementedException();
-    }
-
-    public ValueTask DisposeAsync()
-    {
-        throw new NotImplementedException();
-    }
-
-    public ValueTask<bool> 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<DeleteSessionResponse> 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
@@ -107,10 +107,8 @@ public async Task Tx()
                     func: async stream =>
                     {
                         var titles = new List<string>();
-                        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,20 +339,15 @@ REPLACE INTO episodes
     private static async Task<List<Series>> ReadSeries(ExecuteQueryStream stream)
     {
         var series = new List<Series>();
-        // 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)
             {
                 series.AddRange(resultSet.Rows.Select(Series.FromRow));
             }
         }
 
-        stream.Response.EnsureSuccess();
-
         return series;
     }
 }