fix: handle exceptions in BulkUpsertImporter DisposeAsync, update tests

Author: LiamHamsters

src/Ydb.Sdk/src/Ado/BulkUpsert/BulkUpsertFormat.cs

@@ -1,13 +1,16 @@
+using Ydb.Sdk.Ado.Internal;
 using Ydb.Sdk.Value;
+using Ydb.Table;
+using Ydb.Table.V1;
 
 namespace Ydb.Sdk.Ado.BulkUpsert;
 
 public sealed class BulkUpsertImporter : IBulkUpsertImporter
 {
     private readonly YdbConnection _connection;
     private readonly string _tablePath;
-    private readonly List<string> _columns;
-    private readonly List<Type> _types;
+    private readonly IReadOnlyList<string> _columns;
+    private readonly IReadOnlyList<Type> _types;
     private readonly int _maxBytes;
     private readonly List<Ydb.Value> _rows = new();
     private bool _disposed;
@@ -17,15 +20,30 @@ public BulkUpsertImporter(
         string tablePath,
         IReadOnlyList<string> columns,
         IReadOnlyList<Type> types,
-        int maxBytes = 1024 * 1024)
+        int maxBytes = 64 * 1024 * 1024
+    )
     {
         _connection = connection;
         _tablePath = tablePath;
-        _columns = columns.ToList();
-        _types = types.ToList();
+        _columns = columns;
+        _types = types;
         _maxBytes = maxBytes;
     }
 
+    public async ValueTask AddRowAsync(params object?[] values)
+    {
+        ThrowIfDisposed();
+        if (values.Length != _columns.Count)
+            throw new ArgumentException("Values count must match columns count", nameof(values));
+
+        var ydbValues = values.Select(v =>
+            v is YdbValue yv ? yv :
+            v is YdbParameter param ? param.YdbValue :
+            throw new ArgumentException("All values must be either YdbValue or YdbParameter")).ToArray();
+
+        await AddRowAsync(ydbValues);
+    }
+
     public async ValueTask AddRowAsync(params YdbValue[] values)
     {
         ThrowIfDisposed();
@@ -43,29 +61,15 @@ public async ValueTask AddRowAsync(params YdbValue[] values)
             await FlushAsync();
     }
 
-    public async ValueTask AddRowAsync(params object?[] values)
-    {
-        ThrowIfDisposed();
-        if (values.Length != _columns.Count)
-            throw new ArgumentException("Values count must match columns count", nameof(values));
-
-        var ydbValues = new YdbValue[values.Length];
-        for (int i = 0; i < values.Length; i++)
-        {
-            ydbValues[i] = YdbValueFromObject(values[i], _types[i]);
-        }
-        await AddRowAsync(ydbValues);
-    }
-
-    public async ValueTask AddRowsAsync(IEnumerable<YdbValue[]> rows, CancellationToken cancellationToken = default)
+    public async ValueTask AddRowsAsync(IEnumerable<object?[]> rows, CancellationToken cancellationToken = default)
     {
         ThrowIfDisposed();
 
         foreach (var values in rows)
             await AddRowAsync(values);
     }
 
-    public async ValueTask AddRowsAsync(IEnumerable<object?[]> rows, CancellationToken cancellationToken = default)
+    public async ValueTask AddRowsAsync(IEnumerable<YdbValue[]> rows, CancellationToken cancellationToken = default)
     {
         ThrowIfDisposed();
 
@@ -78,7 +82,22 @@ public async ValueTask FlushAsync(CancellationToken cancellationToken = default)
         ThrowIfDisposed();
         if (_rows.Count == 0) return;
 
-        await _connection.BulkUpsertProtoAsync(_tablePath, GetStructType(), _rows.ToList(), cancellationToken);
+        var listValue = new Ydb.Value();
+        listValue.Items.AddRange(_rows);
+
+        var typedValue = new TypedValue { Type = GetStructType(), Value = listValue };
+        var req = new BulkUpsertRequest { Table = _tablePath, Rows = typedValue };
+
+        var resp = await _connection.Session.Driver.UnaryCall(
+            TableService.BulkUpsertMethod,
+            req,
+            new GrpcRequestSettings { CancellationToken = cancellationToken }
+        ).ConfigureAwait(false);
+
+        var operation = resp.Operation;
+        if (operation.Status.IsNotSuccess())
+            throw YdbException.FromServer(operation.Status, operation.Issues);
+
         _rows.Clear();
     }
 
@@ -104,23 +123,4 @@ private void ThrowIfDisposed()
         if (_disposed)
             throw new ObjectDisposedException(nameof(BulkUpsertImporter));
     }
-
-    private static YdbValue YdbValueFromObject(object? value, Type columnType)
-    {
-        switch (value)
-        {
-            case YdbValue ydbValue:
-                return ydbValue;
-            default:
-                switch (columnType.TypeId)
-                {
-                    case Type.Types.PrimitiveTypeId.Int32:
-                        return YdbValue.MakeInt32(Convert.ToInt32(value));
-                    case Type.Types.PrimitiveTypeId.Utf8:
-                        return YdbValue.MakeUtf8(value?.ToString()!);
-                    default:
-                        throw new NotSupportedException($"Type '{columnType.TypeId}' not supported in YdbValueFromObject");
-                }
-        }
-    }
 }

src/Ydb.Sdk/src/Ado/BulkUpsert/YdbBulkUpsertImporter.cs

@@ -13,6 +13,8 @@ namespace Ydb.Sdk.Ado;
 
 public sealed class YdbConnection : DbConnection
 {
+    private const int MaxSendBufferSize = 64 * 1024 * 1024;
+    
     private static readonly StateChangeEventArgs ClosedToOpenEventArgs =
         new(ConnectionState.Closed, ConnectionState.Open);
 
@@ -56,45 +58,20 @@ public YdbConnection(YdbConnectionStringBuilder connectionStringBuilder)
         ConnectionStringBuilder = connectionStringBuilder;
     }
 
-    public async Task BulkUpsertProtoAsync(
-        string tablePath,
-        Type structType,
-        IReadOnlyList<Ydb.Value> chunk,
-        CancellationToken cancellationToken = default)
-    {
-        if (CurrentTransaction is { Completed: false })
-            throw new InvalidOperationException("BulkUpsertProto cannot be used inside an active transaction.");
-
-        var listValue = new Ydb.Value();
-        listValue.Items.AddRange(chunk);
-
-        var typedValue = new TypedValue { Type = structType, Value = listValue };
-        var req = new BulkUpsertRequest { Table = tablePath, Rows = typedValue };
-
-        var resp = await Session.Driver.UnaryCall(
-            TableService.BulkUpsertMethod,
-            req,
-            new GrpcRequestSettings { CancellationToken = cancellationToken }
-        ).ConfigureAwait(false);
-
-        var operation = resp.Operation;
-        if (operation.Status.IsNotSuccess())
-            throw YdbException.FromServer(operation.Status, operation.Issues);
-    }
-    
     public IBulkUpsertImporter BeginBulkUpsertImport(
-        string tablePath,
+        string name,
         IReadOnlyList<string> columns,
-        IReadOnlyList<Type> types,
-        int maxBytes = 1024 * 1024)
+        IReadOnlyList<Type> types)
     {
         ThrowIfConnectionClosed();
         if (CurrentTransaction is { Completed: false })
             throw new InvalidOperationException("BulkUpsert cannot be used inside an active transaction.");
 
-        return new BulkUpsertImporter(this, tablePath, columns, types, maxBytes);
-    }
+        var database = (ConnectionStringBuilder.Database ?? "").TrimEnd('/');
+        var tablePath = string.IsNullOrEmpty(database) ? name : $"{database}/{name}";
 
+        return new BulkUpsertImporter(this, tablePath, columns, types, MaxSendBufferSize);
+    }
 
     protected override YdbTransaction BeginDbTransaction(IsolationLevel isolationLevel)
     {

src/Ydb.Sdk/src/Ado/YdbConnection.cs

@@ -313,8 +313,6 @@ protected override async Task OnDisposeAsync() =>
     public async Task BulkUpsertImporter_HappyPath_AddRows_Flush_Dispose()
     {
         var tableName = $"BulkImporter_{Guid.NewGuid():N}";
-        var database = new YdbConnectionStringBuilder(_connectionStringTls).Database.TrimEnd('/');
-        var absTablePath = string.IsNullOrEmpty(database) ? tableName : $"{database}/{tableName}";
 
         var conn = new YdbConnection(_connectionStringTls);
         await conn.OpenAsync();
@@ -337,7 +335,7 @@ PRIMARY KEY (Id)
                 new Type { TypeId = Type.Types.PrimitiveTypeId.Utf8 }
             };
 
-            await using (var importer = conn.BeginBulkUpsertImport(absTablePath, columns, types, maxBytes: 1024))
+            await using (var importer = conn.BeginBulkUpsertImport(tableName, columns, types))
             {
                 await importer.AddRowAsync(YdbValue.MakeInt32(1), YdbValue.MakeUtf8("Alice"));
                 await importer.AddRowAsync(YdbValue.MakeInt32(2), YdbValue.MakeUtf8("Bob"));
@@ -353,7 +351,7 @@ PRIMARY KEY (Id)
                 Assert.Equal(2, count);
             }
 
-            await using (var importer = conn.BeginBulkUpsertImport(absTablePath, columns, types, maxBytes: 1))
+            await using (var importer = conn.BeginBulkUpsertImport(tableName, columns, types))
             {
                 var rows = new List<YdbValue[]>
                 {
@@ -409,7 +407,7 @@ PRIMARY KEY (Id)
                 new Type { TypeId = Type.Types.PrimitiveTypeId.Utf8 }
             };
 
-            await using (var importer = conn.BeginBulkUpsertImport(tableName, columns, types, maxBytes: 1024))
+            await using (var importer = conn.BeginBulkUpsertImport(tableName, columns, types))
             {
                 await Assert.ThrowsAsync<ArgumentException>(async () =>
                 {
@@ -438,8 +436,6 @@ await Assert.ThrowsAsync<ArgumentException>(async () =>
     public async Task BulkUpsertImporter_DisposeIsIdempotent()
     {
         var tableName = $"BulkImporter_{Guid.NewGuid():N}";
-        var database = new YdbConnectionStringBuilder(_connectionStringTls).Database.TrimEnd('/');
-        var absTablePath = string.IsNullOrEmpty(database) ? tableName : $"{database}/{tableName}";
 
         var conn = new YdbConnection(_connectionStringTls);
         await conn.OpenAsync();
@@ -462,7 +458,7 @@ PRIMARY KEY (Id)
                 new Type { TypeId = Type.Types.PrimitiveTypeId.Utf8 }
             };
 
-            var importer = conn.BeginBulkUpsertImport(absTablePath, columns, types, maxBytes: 1024);
+            var importer = conn.BeginBulkUpsertImport(tableName, columns, types);
             await importer.AddRowAsync(YdbValue.MakeInt32(1), YdbValue.MakeUtf8("A"));
             await importer.DisposeAsync();
 
@@ -475,15 +471,12 @@ PRIMARY KEY (Id)
             await dropCmd.ExecuteNonQueryAsync();
         }
     }
-    
+
     [Fact]
     public async Task BulkUpsertImporter_MultipleImporters_Parallel()
     {
         var table1 = $"BulkImporter_{Guid.NewGuid():N}_1";
         var table2 = $"BulkImporter_{Guid.NewGuid():N}_2";
-        var database = new YdbConnectionStringBuilder(_connectionStringTls).Database.TrimEnd('/');
-        var absTablePath1 = string.IsNullOrEmpty(database) ? table1 : $"{database}/{table1}";
-        var absTablePath2 = string.IsNullOrEmpty(database) ? table2 : $"{database}/{table2}";
 
         var conn = new YdbConnection(_connectionStringTls);
         await conn.OpenAsync();
@@ -510,13 +503,13 @@ PRIMARY KEY (Id)
             await Task.WhenAll(
                 Task.Run(async () =>
                 {
-                    await using var importer = conn.BeginBulkUpsertImport(absTablePath1, columns, types, maxBytes: 1024);
+                    await using var importer = conn.BeginBulkUpsertImport(table1, columns, types);
                     for (int i = 0; i < 20; i++)
                         await importer.AddRowAsync(YdbValue.MakeInt32(i), YdbValue.MakeUtf8($"A{i}"));
                 }),
                 Task.Run(async () =>
                 {
-                    await using var importer = conn.BeginBulkUpsertImport(absTablePath2, columns, types, maxBytes: 1024);
+                    await using var importer = conn.BeginBulkUpsertImport(table2, columns, types);
                     for (int i = 0; i < 20; i++)
                         await importer.AddRowAsync(YdbValue.MakeInt32(i), YdbValue.MakeUtf8($"B{i}"));
                 })
@@ -540,55 +533,36 @@ await Task.WhenAll(
             }
         }
     }
-    
+
     [Fact]
-    public async Task BulkUpsertImporter_TinyMaxBytes_FlushesEveryRow()
+    public async Task BulkUpsertImporter_ThrowsOnNonexistentTable()
     {
-        var tableName = $"BulkImporter_{Guid.NewGuid():N}";
-        var database = new YdbConnectionStringBuilder(_connectionStringTls).Database.TrimEnd('/');
-        var absTablePath = string.IsNullOrEmpty(database) ? tableName : $"{database}/{tableName}";
-
+        var tableName = $"Nonexistent_{Guid.NewGuid():N}";
         var conn = new YdbConnection(_connectionStringTls);
         await conn.OpenAsync();
-        try
+
+        var columns = new[] { "Id", "Name" };
+        var types = new[]
         {
-            await using (var createCmd = conn.CreateCommand())
-            {
-                createCmd.CommandText = $@"
-            CREATE TABLE {tableName} (
-                Id Int32,
-                Name Utf8,
-                PRIMARY KEY (Id)
-            )";
-                await createCmd.ExecuteNonQueryAsync();
-            }
+            new Type { TypeId = Type.Types.PrimitiveTypeId.Int32 },
+            new Type { TypeId = Type.Types.PrimitiveTypeId.Utf8 }
+        };
 
-            var columns = new[] { "Id", "Name" };
-            var types = new[]
-            {
-                new Type { TypeId = Type.Types.PrimitiveTypeId.Int32 },
-                new Type { TypeId = Type.Types.PrimitiveTypeId.Utf8 }
-            };
+        var importer = conn.BeginBulkUpsertImport(tableName, columns, types);
 
-            await using (var importer = conn.BeginBulkUpsertImport(absTablePath, columns, types, maxBytes: 1))
-            {
-                for (int i = 0; i < 10; i++)
-                    await importer.AddRowAsync(YdbValue.MakeInt32(i), YdbValue.MakeUtf8($"Name{i}"));
-                Assert.Empty(importer.GetBufferedRows());
-            }
+        await Assert.ThrowsAsync<YdbException>(async () =>
+        {
+            await importer.AddRowAsync(YdbValue.MakeInt32(1), YdbValue.MakeUtf8("NotExists"));
+            await importer.FlushAsync();
+        });
 
-            await using (var checkCmd = conn.CreateCommand())
-            {
-                checkCmd.CommandText = $"SELECT COUNT(*) FROM {tableName}";
-                var count = Convert.ToInt32(await checkCmd.ExecuteScalarAsync());
-                Assert.Equal(10, count);
-            }
+        try
+        {
+            await importer.DisposeAsync();
         }
-        finally
+        catch (YdbException)
         {
-            await using var dropCmd = conn.CreateCommand();
-            dropCmd.CommandText = $"DROP TABLE {tableName}";
-            await dropCmd.ExecuteNonQueryAsync();
+            // success
         }
     }
 }