chore: add GetTotalUpdateCount method to .NET (#701)

* chore: add GetTotalUpdateCount method to .NET

Add a GetTotalUpdateCount method to the .NET wrapper that can be used to
return the total update count of a set of SQL statements. This will be used
by the ADO.NET ExecuteNonQuery method.

See https://learn.microsoft.com/en-us/dotnet/api/system.data.sqlclient.sqlcommand.executenonquery

* fix: only return -1 if there are no DML statements

Author: olavloite

spannerlib/wrappers/spannerlib-dotnet/spannerlib-dotnet-tests/RowsTests.cs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+using System.Text;
 using Google.Cloud.Spanner.Admin.Database.V1;
 using Google.Cloud.Spanner.V1;
 using Google.Cloud.SpannerLib.MockServer;
@@ -429,6 +430,51 @@ public async Task TestMultipleMixedStatements([Values] LibType libType, [Values(
         Assert.That(numResultSets, Is.EqualTo(5));
     }
 
+    [Test]
+    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)
+    {
+        var updateCount = 3L;
+        var dml = "update my_table set value=1 where id in (1,2,3)";
+        Fixture.SpannerMock.AddOrUpdateStatementResult(dml, StatementResult.CreateUpdateCount(updateCount));
+        Fixture.SpannerMock.AddOrUpdateStatementResult(dml + ";", StatementResult.CreateUpdateCount(updateCount));
+        
+        var rowType = RandomResultSetGenerator.GenerateAllTypesRowType();
+        var results = RandomResultSetGenerator.Generate(rowType, numRows);
+        var query = "select * from random";
+        Fixture.SpannerMock.AddOrUpdateStatementResult(query, StatementResult.CreateQuery(results));
+
+        // Create a SQL string containing a mix of DML and queries.
+        var builder = new StringBuilder();
+        for (var i = 0; i < numDmlStatements; i++)
+        {
+            while (Random.Shared.Next(2) == 0)
+            {
+                builder.Append(query).Append(';');
+            }
+            builder.Append(dml).Append(';');
+            while (Random.Shared.Next(5) == 0)
+            {
+                builder.Append(query).Append(';');
+            }
+        }
+        var sql = builder.ToString();
+        if (string.IsNullOrEmpty(sql))
+        {
+            sql = query;
+        }
+
+        await using var pool = Pool.Create(SpannerLibDictionary[libType], ConnectionString);
+        await using var connection = pool.CreateConnection();
+        await using var rows = async
+            ? await connection.ExecuteAsync(new ExecuteSqlRequest { Sql = sql }, prefetchRows)
+            // ReSharper disable once MethodHasAsyncOverload
+            : connection.Execute(new ExecuteSqlRequest { Sql = sql }, prefetchRows);
+
+        // ReSharper disable once MethodHasAsyncOverload
+        var totalUpdateCount = async ? await rows.GetTotalUpdateCountAsync() : rows.GetTotalUpdateCount();
+        Assert.That(totalUpdateCount, Is.EqualTo(numDmlStatements == 0 ? -1 : updateCount * numDmlStatements));
+    }
+    
     [Test]
     public async Task TestMultipleMixedStatementsWithErrors(
         [Values] LibType libType,

spannerlib/wrappers/spannerlib-dotnet/spannerlib-dotnet/Rows.cs

@@ -62,6 +62,8 @@ public long UpdateCount
             return -1L;
         }
     }
+    
+    private bool _hasReadAllResults;
 
     public Rows(Connection connection, long id, bool initMetadata = true) : base(connection.Spanner, id)
     {
@@ -97,12 +99,78 @@ public Rows(Connection connection, long id, bool initMetadata = true) : base(con
         return await Spanner.NextAsync(this, 1, ISpannerLib.RowEncoding.Proto, cancellationToken).ConfigureAwait(false);
     }
 
+    /// <summary>
+    /// Gets the total update count in this Rows object. This consumes all data in all the result sets.
+    /// This method should only be called if the caller is only interested in the update count, and not in any of the
+    /// rows in the result sets.
+    /// </summary>
+    /// <returns>
+    /// The total update count of all the result sets in this Rows object.
+    /// If the SQL string only contained non-DML statements, then the update count is -1. If the SQL string contained at
+    /// least one DML statement, then the returned value is the sum of all update counts of the DML statements in the
+    /// SQL string that was executed.
+    /// </returns>
+    public long GetTotalUpdateCount()
+    {
+        long result = -1;
+        var hasUpdateCount = false;
+        do
+        {
+            var updateCount = UpdateCount;
+            if (updateCount > -1)
+            {
+                if (!hasUpdateCount)
+                {
+                    hasUpdateCount = true;
+                    result = 0;
+                }
+                result += updateCount;
+            }
+        } while (NextResultSet());
+        return result;
+    }
+
+    /// <summary>
+    /// Gets the total update count in this Rows object. This consumes all data in all the result sets.
+    /// This method should only be called if the caller is only interested in the update count, and not in any of the
+    /// rows in the result sets.
+    /// </summary>
+    /// <returns>
+    /// The total update count of all the result sets in this Rows object.
+    /// If the SQL string only contained non-DML statements, then the update count is -1. If the SQL string contained at
+    /// least one DML statement, then the returned value is the sum of all update counts of the DML statements in the
+    /// SQL string that was executed.
+    /// </returns>
+    public async Task<long> GetTotalUpdateCountAsync(CancellationToken cancellationToken = default)
+    {
+        long result = -1;
+        var hasUpdateCount = false;
+        do
+        {
+            var updateCount = UpdateCount;
+            if (updateCount > -1)
+            {
+                if (!hasUpdateCount)
+                {
+                    hasUpdateCount = true;
+                    result = 0;
+                }
+                result += updateCount;
+            }
+        } while (await NextResultSetAsync(cancellationToken).ConfigureAwait(false));
+        return result;
+    }
+
     /// <summary>
     /// Moves the cursor to the next result set in this Rows object.
     /// </summary>
     /// <returns>True if there was another result set, and false otherwise</returns>
     public virtual bool NextResultSet()
     {
+        if (_hasReadAllResults)
+        {
+            return false;
+        }
         return NextResultSet(Spanner.NextResultSet(this));
     }
 
@@ -112,13 +180,18 @@ public virtual bool NextResultSet()
     /// <returns>True if there was another result set, and false otherwise</returns>
     public virtual async Task<bool> NextResultSetAsync(CancellationToken cancellationToken = default)
     {
+        if (_hasReadAllResults)
+        {
+            return false;
+        }
         return NextResultSet(await Spanner.NextResultSetAsync(this, cancellationToken).ConfigureAwait(false));
     }
 
     private bool NextResultSet(ResultSetMetadata? metadata)
     {
         if (metadata == null)
         {
+            _hasReadAllResults = true;
             return false;
         }
         _metadata = metadata;