|
12 | 12 | // See the License for the specific language governing permissions and |
13 | 13 | // limitations under the License. |
14 | 14 |
|
| 15 | +using System.Text; |
15 | 16 | using Google.Cloud.Spanner.Admin.Database.V1; |
16 | 17 | using Google.Cloud.Spanner.V1; |
17 | 18 | using Google.Cloud.SpannerLib.MockServer; |
@@ -83,6 +84,50 @@ public void TestRandomResults([Values] LibType libType, [Values(0, 1, 10)] int n |
83 | 84 | var request = Fixture.SpannerMock.Requests.OfType<ExecuteSqlRequest>().First(); |
84 | 85 | Assert.That(request.Transaction?.SingleUse?.ReadOnly?.HasStrong ?? false); |
85 | 86 | } |
| 87 | + |
| 88 | + [Test] |
| 89 | + public async Task TestCancellation([Values] LibType libType, [Values(0, 1, 10)] int numRows, [Values(0, 1, 5, 9, 10, 11)] int prefetchRows) |
| 90 | + { |
| 91 | + var rowType = RandomResultSetGenerator.GenerateAllTypesRowType(); |
| 92 | + var results = RandomResultSetGenerator.Generate(rowType, numRows); |
| 93 | + Fixture.SpannerMock.AddOrUpdateStatementResult("select * from random", StatementResult.CreateQuery(results)); |
| 94 | + |
| 95 | + CancellationTokenSource cts = new CancellationTokenSource(); |
| 96 | + var numReadRowsBeforeCancellation = Random.Shared.Next(numRows); |
| 97 | + await using var pool = Pool.Create(SpannerLibDictionary[libType], ConnectionString); |
| 98 | + await using var connection = pool.CreateConnection(); |
| 99 | + if (numReadRowsBeforeCancellation == 0) |
| 100 | + { |
| 101 | + await cts.CancelAsync(); |
| 102 | + var exception = Assert.CatchAsync<Exception>(async () => await connection.ExecuteAsync(new ExecuteSqlRequest { Sql = "select * from random" }, prefetchRows, cts.Token)); |
| 103 | + Assert.That(exception, Is.Not.Null); |
| 104 | + Assert.That(exception, Is.InstanceOf<OperationCanceledException>().Or.InstanceOf<RpcException>().Or.InstanceOf<TaskCanceledException>()); |
| 105 | + } |
| 106 | + else |
| 107 | + { |
| 108 | + await using var rows = await connection.ExecuteAsync(new ExecuteSqlRequest { Sql = "select * from random" }, |
| 109 | + prefetchRows, cts.Token); |
| 110 | + |
| 111 | + var rowCount = 0; |
| 112 | + while (await rows.NextAsync(cts.Token) is not null) |
| 113 | + { |
| 114 | + rowCount++; |
| 115 | + if (numReadRowsBeforeCancellation == rowCount) |
| 116 | + { |
| 117 | + await cts.CancelAsync(); |
| 118 | + var exception = Assert.CatchAsync<Exception>(async () => await rows.NextAsync(cts.Token)); |
| 119 | + Assert.That(exception, Is.Not.Null); |
| 120 | + Assert.That(exception, Is.InstanceOf<OperationCanceledException>().Or.InstanceOf<RpcException>().Or.InstanceOf<TaskCanceledException>()); |
| 121 | + break; |
| 122 | + } |
| 123 | + } |
| 124 | + Assert.That(rowCount, Is.LessThan(numRows)); |
| 125 | + |
| 126 | + Assert.That(Fixture.SpannerMock.Requests.OfType<ExecuteSqlRequest>().Count(), Is.EqualTo(1)); |
| 127 | + var request = Fixture.SpannerMock.Requests.OfType<ExecuteSqlRequest>().First(); |
| 128 | + Assert.That(request.Transaction?.SingleUse?.ReadOnly?.HasStrong ?? false); |
| 129 | + } |
| 130 | + } |
86 | 131 |
|
87 | 132 | [Test] |
88 | 133 | public void TestLargeStringValue([Values] LibType libType) |
@@ -385,6 +430,51 @@ public async Task TestMultipleMixedStatements([Values] LibType libType, [Values( |
385 | 430 | Assert.That(numResultSets, Is.EqualTo(5)); |
386 | 431 | } |
387 | 432 |
|
| 433 | + [Test] |
| 434 | + public async Task TestGetAllUpdateCounts([Values] LibType libType, [Values(0, 1, 2, 5)] int numDmlStatements, [Values(2, 10)] int numRows, [Values(0, 1, 2, 3, 5, 9, 10, 11)] int prefetchRows, [Values] bool async) |
| 435 | + { |
| 436 | + var updateCount = 3L; |
| 437 | + var dml = "update my_table set value=1 where id in (1,2,3)"; |
| 438 | + Fixture.SpannerMock.AddOrUpdateStatementResult(dml, StatementResult.CreateUpdateCount(updateCount)); |
| 439 | + Fixture.SpannerMock.AddOrUpdateStatementResult(dml + ";", StatementResult.CreateUpdateCount(updateCount)); |
| 440 | + |
| 441 | + var rowType = RandomResultSetGenerator.GenerateAllTypesRowType(); |
| 442 | + var results = RandomResultSetGenerator.Generate(rowType, numRows); |
| 443 | + var query = "select * from random"; |
| 444 | + Fixture.SpannerMock.AddOrUpdateStatementResult(query, StatementResult.CreateQuery(results)); |
| 445 | + |
| 446 | + // Create a SQL string containing a mix of DML and queries. |
| 447 | + var builder = new StringBuilder(); |
| 448 | + for (var i = 0; i < numDmlStatements; i++) |
| 449 | + { |
| 450 | + while (Random.Shared.Next(2) == 0) |
| 451 | + { |
| 452 | + builder.Append(query).Append(';'); |
| 453 | + } |
| 454 | + builder.Append(dml).Append(';'); |
| 455 | + while (Random.Shared.Next(5) == 0) |
| 456 | + { |
| 457 | + builder.Append(query).Append(';'); |
| 458 | + } |
| 459 | + } |
| 460 | + var sql = builder.ToString(); |
| 461 | + if (string.IsNullOrEmpty(sql)) |
| 462 | + { |
| 463 | + sql = query; |
| 464 | + } |
| 465 | + |
| 466 | + await using var pool = Pool.Create(SpannerLibDictionary[libType], ConnectionString); |
| 467 | + await using var connection = pool.CreateConnection(); |
| 468 | + await using var rows = async |
| 469 | + ? await connection.ExecuteAsync(new ExecuteSqlRequest { Sql = sql }, prefetchRows) |
| 470 | + // ReSharper disable once MethodHasAsyncOverload |
| 471 | + : connection.Execute(new ExecuteSqlRequest { Sql = sql }, prefetchRows); |
| 472 | + |
| 473 | + // ReSharper disable once MethodHasAsyncOverload |
| 474 | + var totalUpdateCount = async ? await rows.GetTotalUpdateCountAsync() : rows.GetTotalUpdateCount(); |
| 475 | + Assert.That(totalUpdateCount, Is.EqualTo(numDmlStatements == 0 ? -1 : updateCount * numDmlStatements)); |
| 476 | + } |
| 477 | + |
388 | 478 | [Test] |
389 | 479 | public async Task TestMultipleMixedStatementsWithErrors( |
390 | 480 | [Values] LibType libType, |
|
0 commit comments