From a565aab49828d1d2f78b4ff2ccd9a2679866ab57 Mon Sep 17 00:00:00 2001
From: XmasApple <y4blokof@yandex-team.ru>
Date: Wed, 15 Nov 2023 18:12:19 +0300
Subject: [PATCH] QueryService: test for idempotency

---
 src/Ydb.Sdk/src/Retry.cs                      |  2 +-
 src/Ydb.Sdk/src/Services/Query/QueryClient.cs | 13 ++--
 .../tests/Query/TestQueryIntegration.cs       | 59 ++++++++++++++++++-
 3 files changed, 67 insertions(+), 7 deletions(-)

diff --git a/src/Ydb.Sdk/src/Retry.cs b/src/Ydb.Sdk/src/Retry.cs
index 30226783..51ab32d7 100644
--- a/src/Ydb.Sdk/src/Retry.cs
+++ b/src/Ydb.Sdk/src/Retry.cs
@@ -47,7 +47,7 @@ public enum Idempotency
     /// <summary> No retry </summary>
     None,
 
-    /// <summary> Retry if IsIdempotent is true </summary>
+    /// <summary> Retry only if IsIdempotent is true </summary>
     Idempotent,
 
     /// <summary> Retry always </summary>
diff --git a/src/Ydb.Sdk/src/Services/Query/QueryClient.cs b/src/Ydb.Sdk/src/Services/Query/QueryClient.cs
index 7c1daecd..c88644a7 100644
--- a/src/Ydb.Sdk/src/Services/Query/QueryClient.cs
+++ b/src/Ydb.Sdk/src/Services/Query/QueryClient.cs
@@ -662,7 +662,7 @@ public async Task<QueryResponseWithResult<T>> DoTx<T>(Func<Tx, Task<T>> func,
                 catch (Exception e)
                 {
                     var status = new Status(
-                        StatusCode.InternalError,
+                        StatusCode.ClientInternalError,
                         $"Failed to execute lambda on tx {tx.TxId}: {e.Message}");
                     var rollbackResponse = await Rollback<T>(session, tx, status);
                     return rollbackResponse;
@@ -671,12 +671,15 @@ public async Task<QueryResponseWithResult<T>> DoTx<T>(Func<Tx, Task<T>> func,
                 var commitResponse = await CommitTransaction(session.Id, tx);
                 if (!commitResponse.Status.IsSuccess)
                 {
-                    return await Rollback<T>(session, tx, commitResponse.Status);
+                    var rollbackResponse = await Rollback<T>(session, tx, commitResponse.Status);
+                    return rollbackResponse;
                 }
 
-                return response is None
-                    ? new QueryResponseWithResult<T>(Status.Success)
-                    : new QueryResponseWithResult<T>(Status.Success, response);
+                return response switch
+                {
+                    None => new QueryResponseWithResult<T>(Status.Success),
+                    _ => new QueryResponseWithResult<T>(Status.Success, response)
+                };
             },
             retrySettings
         );
diff --git a/src/Ydb.Sdk/tests/Query/TestQueryIntegration.cs b/src/Ydb.Sdk/tests/Query/TestQueryIntegration.cs
index bd052df5..f86269d5 100644
--- a/src/Ydb.Sdk/tests/Query/TestQueryIntegration.cs
+++ b/src/Ydb.Sdk/tests/Query/TestQueryIntegration.cs
@@ -1,6 +1,7 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 using Xunit;
+using Ydb.Sdk.Client;
 using Ydb.Sdk.Services.Query;
 using Ydb.Sdk.Services.Table;
 using Ydb.Sdk.Value;
@@ -43,7 +44,7 @@ public async Task TestSimpleSelect()
         using var client = new QueryClient(driver);
 
         const string queryString = "SELECT 2 + 3 AS sum";
-        
+
         var responseQuery = await client.Query(queryString, async stream =>
         {
             var rows = new List<Sdk.Value.ResultSet.Row>();
@@ -177,4 +178,60 @@ DELETE FROM {tableName}
 
         await DropTable(tableClient, tableName);
     }
+
+    [Fact]
+    public async Task TestDoTxRollback()
+    {
+        await using var driver = await Driver.CreateInitialized(_driverConfig, _loggerFactory);
+        using var client = new QueryClient(driver);
+
+        var response = await client.DoTx(_ =>
+        {
+            var response = new ClientInternalErrorResponse("test rollback if status unsuccessful");
+            response.EnsureSuccess();
+            return Task.CompletedTask;
+        });
+        Assert.Equal(StatusCode.ClientInternalError, response.Status.StatusCode);
+
+
+        response = await client.DoTx(_ => throw new ArithmeticException("2 + 2 = 5"));
+        Assert.Equal(StatusCode.ClientInternalError, response.Status.StatusCode);
+    }
+
+    [Theory]
+    [InlineData(StatusCode.ClientInternalError, StatusCode.Success, 2, true)]
+    [InlineData(StatusCode.ClientInternalError, StatusCode.ClientInternalError, 1, false)]
+    [InlineData(StatusCode.InternalError, StatusCode.InternalError, 1, true)]
+    [InlineData(StatusCode.Aborted, StatusCode.Success, 2, false)]
+    public async Task TestIdempotency(StatusCode statusCode, StatusCode expectedStatusCode, int expectedAttempts,
+        bool isIdempotent)
+    {
+        await using var driver = await Driver.CreateInitialized(_driverConfig, _loggerFactory);
+        using var client = new QueryClient(driver);
+
+        var attempts = 0;
+        var response = await client.Query("SELECT 1", async stream =>
+            {
+                attempts += 1;
+                var rows = new List<Sdk.Value.ResultSet.Row>();
+                await foreach (var part in stream)
+                {
+                    if (part.ResultSet is not null)
+                    {
+                        rows.AddRange(part.ResultSet.Rows);
+                    }
+                }
+
+                if (attempts == 1)
+                {
+                    throw new StatusUnsuccessfulException(new Status(statusCode, "test idempotency"));
+                }
+
+                return rows;
+            },
+            retrySettings: new RetrySettings { IsIdempotent = isIdempotent });
+
+        Assert.Equal(expectedStatusCode, response.Status.StatusCode);
+        Assert.Equal(expectedAttempts, attempts);
+    }
 }