bulk upsert: add object[] support, simplify implementation

LiamHamsters · LiamHamsters · commit 1879bca17fcf · 2025-07-29T15:02:49.000+03:00
diff --git a/src/Ydb.Sdk/src/Ado/BulkUpsert/IBulkUpsertImporter.cs b/src/Ydb.Sdk/src/Ado/BulkUpsert/IBulkUpsertImporter.cs
@@ -1,7 +1,15 @@
+using Ydb.Sdk.Value;
+
 namespace Ydb.Sdk.Ado.BulkUpsert;
 
 public interface IBulkUpsertImporter : IAsyncDisposable
 {
-    ValueTask AddRowAsync(params Ydb.Value[] values);
+    ValueTask AddRowAsync(params YdbValue[] values);
+    ValueTask AddRowAsync(params object?[] values);
+
+    ValueTask AddRowsAsync(IEnumerable<YdbValue[]> rows, CancellationToken cancellationToken = default);
+    ValueTask AddRowsAsync(IEnumerable<object?[]> rows, CancellationToken cancellationToken = default);
+
     ValueTask FlushAsync(CancellationToken cancellationToken = default);
-}
+    IReadOnlyList<Ydb.Value> GetBufferedRows();
+}
diff --git a/src/Ydb.Sdk/src/Ado/BulkUpsert/YdbBulkUpsertImporter.cs b/src/Ydb.Sdk/src/Ado/BulkUpsert/YdbBulkUpsertImporter.cs
@@ -2,18 +2,17 @@
 
 namespace Ydb.Sdk.Ado.BulkUpsert;
 
-public sealed class YdbBulkUpsertProtoImporter : IAsyncDisposable
+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 int _maxBytes;
-
     private readonly List<Ydb.Value> _rows = new();
     private bool _disposed;
 
-    public YdbBulkUpsertProtoImporter(
+    public BulkUpsertImporter(
         YdbConnection connection,
         string tablePath,
         IReadOnlyList<string> columns,
@@ -22,15 +21,14 @@ public YdbBulkUpsertProtoImporter(
     {
         _connection = connection;
         _tablePath = tablePath;
-        _columns = columns is List<string> colList ? colList : new List<string>(columns);
-        _types = types is List<Type> typList ? typList : new List<Type>(types);
+        _columns = columns.ToList();
+        _types = types.ToList();
         _maxBytes = maxBytes;
     }
 
     public async ValueTask AddRowAsync(params YdbValue[] values)
     {
         ThrowIfDisposed();
-
         if (values.Length != _columns.Count)
             throw new ArgumentException("Values count must match columns count", nameof(values));
 
@@ -41,76 +39,88 @@ public async ValueTask AddRowAsync(params YdbValue[] values)
         _rows.Add(structValue);
 
         var totalSize = _rows.Sum(r => r.CalculateSize());
-
         if (totalSize >= _maxBytes)
             await FlushAsync();
     }
 
-    public async ValueTask AddRowsAsync(IEnumerable<YdbValue[]> rows, CancellationToken cancellationToken = default)
+    public async ValueTask AddRowAsync(params object?[] values)
     {
         ThrowIfDisposed();
+        if (values.Length != _columns.Count)
+            throw new ArgumentException("Values count must match columns count", nameof(values));
 
-        foreach (var values in rows)
+        var ydbValues = new YdbValue[values.Length];
+        for (int i = 0; i < values.Length; i++)
         {
-            if (values.Length != _columns.Count)
-                throw new ArgumentException("Values count must match columns count", nameof(values));
-
-            var dict = _columns.Zip(values, (name, value) => new KeyValuePair<string, YdbValue>(name, value))
-                .ToDictionary(x => x.Key, x => x.Value);
-
-            var structValue = YdbValue.MakeStruct(dict).GetProto().Value;
-            _rows.Add(structValue);
+            ydbValues[i] = YdbValueFromObject(values[i], _types[i]);
+        }
+        await AddRowAsync(ydbValues);
+    }
 
-            var totalSize = _rows.Sum(r => r.CalculateSize());
+    public async ValueTask AddRowsAsync(IEnumerable<YdbValue[]> rows, CancellationToken cancellationToken = default)
+    {
+        ThrowIfDisposed();
 
-            if (totalSize >= _maxBytes)
-                await FlushAsync(cancellationToken);
-        }
+        foreach (var values in rows)
+            await AddRowAsync(values);
     }
 
-    public async ValueTask DisposeAsync()
+    public async ValueTask AddRowsAsync(IEnumerable<object?[]> rows, CancellationToken cancellationToken = default)
     {
-        if (_disposed) return;
-        await FlushAsync();
-        _disposed = true;
+        ThrowIfDisposed();
+
+        foreach (var values in rows)
+            await AddRowAsync(values);
     }
 
     public async ValueTask FlushAsync(CancellationToken cancellationToken = default)
     {
         ThrowIfDisposed();
+        if (_rows.Count == 0) return;
 
-        if (_rows.Count == 0)
-            return;
+        await _connection.BulkUpsertProtoAsync(_tablePath, GetStructType(), _rows.ToList(), cancellationToken);
+        _rows.Clear();
+    }
 
-        await _connection.BulkUpsertProtoAsync(
-            _tablePath,
-            GetStructType(),
-            new List<Ydb.Value>(_rows),
-            cancellationToken);
+    public IReadOnlyList<Ydb.Value> GetBufferedRows() => _rows;
 
-        _rows.Clear();
+    public async ValueTask DisposeAsync()
+    {
+        if (_disposed) return;
+        await FlushAsync();
+        _disposed = true;
     }
 
     private Type GetStructType()
     {
         var structType = new Type { StructType = new StructType() };
         for (var i = 0; i < _columns.Count; i++)
-        {
-            structType.StructType.Members.Add(new StructMember
-            {
-                Name = _columns[i],
-                Type = _types[i]
-            });
-        }
-
+            structType.StructType.Members.Add(new StructMember { Name = _columns[i], Type = _types[i] });
         return structType;
     }
 
     private void ThrowIfDisposed()
     {
         if (_disposed)
-            throw new ObjectDisposedException(nameof(YdbBulkUpsertProtoImporter));
+            throw new ObjectDisposedException(nameof(BulkUpsertImporter));
     }
 
-    public IReadOnlyList<Ydb.Value> GetBufferedRows() => _rows;
-}
+    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");
+                }
+        }
+    }
+}
diff --git a/src/Ydb.Sdk/src/Ado/YdbConnection.cs b/src/Ydb.Sdk/src/Ado/YdbConnection.cs
@@ -1,10 +1,10 @@
 using System.Data;
 using System.Data.Common;
 using System.Diagnostics.CodeAnalysis;
+using Ydb.Sdk.Ado.BulkUpsert;
 using Ydb.Sdk.Ado.Internal;
 using Ydb.Sdk.Ado.Session;
 using Ydb.Sdk.Services.Query;
-using Ydb.Sdk.Value;
 using Ydb.Table;
 using Ydb.Table.V1;
 using static System.Data.IsolationLevel;
@@ -56,38 +56,20 @@ public YdbConnection(YdbConnectionStringBuilder connectionStringBuilder)
         ConnectionStringBuilder = connectionStringBuilder;
     }
 
-    public async Task BulkUpsertAsync(
+    public async Task BulkUpsertProtoAsync(
         string tablePath,
-        IReadOnlyList<string> columns,
-        IReadOnlyList<IReadOnlyList<object?>> rows,
+        Type structType,
+        IReadOnlyList<Ydb.Value> chunk,
         CancellationToken cancellationToken = default)
     {
         if (CurrentTransaction is { Completed: false })
-            throw new InvalidOperationException("BulkUpsert cannot be used inside an active transaction.");
+            throw new InvalidOperationException("BulkUpsertProto cannot be used inside an active transaction.");
 
-        if (columns == null || columns.Count == 0)
-            throw new ArgumentException("Columns must not be empty", nameof(columns));
-        if (rows == null || rows.Count == 0)
-            throw new ArgumentException("Rows collection is empty", nameof(rows));
+        var listValue = new Ydb.Value();
+        listValue.Items.AddRange(chunk);
 
-        var structs = rows.Select(row =>
-        {
-            if (row.Count != columns.Count)
-                throw new ArgumentException("Each row must have the same number of elements as columns");
-            var members = columns
-                .Select((col, i) =>
-                    new KeyValuePair<string, YdbValue>(col, new YdbParameter { Value = row[i] }.YdbValue))
-                .ToDictionary(x => x.Key, x => x.Value);
-            return YdbValue.MakeStruct(members);
-        }).ToList();
-
-        var list = YdbValue.MakeList(structs);
-
-        var req = new BulkUpsertRequest
-        {
-            Table = tablePath,
-            Rows = list.GetProto()
-        };
+        var typedValue = new TypedValue { Type = structType, Value = listValue };
+        var req = new BulkUpsertRequest { Table = tablePath, Rows = typedValue };
 
         var resp = await Session.Driver.UnaryCall(
             TableService.BulkUpsertMethod,
@@ -97,51 +79,23 @@ public async Task BulkUpsertAsync(
 
         var operation = resp.Operation;
         if (operation.Status.IsNotSuccess())
-        {
             throw YdbException.FromServer(operation.Status, operation.Issues);
-        }
     }
-
-    public async Task BulkUpsertProtoAsync(
+    
+    public IBulkUpsertImporter BeginBulkUpsertImport(
         string tablePath,
-        Type structType,
-        IReadOnlyList<Ydb.Value> chunk,
-        CancellationToken cancellationToken = default,
-        int retryCount = 3)
+        IReadOnlyList<string> columns,
+        IReadOnlyList<Type> types,
+        int maxBytes = 1024 * 1024)
     {
+        ThrowIfConnectionClosed();
         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 };
+            throw new InvalidOperationException("BulkUpsert cannot be used inside an active transaction.");
 
-        var attempt = 0;
-        while (true)
-        {
-            try
-            {
-                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);
-                return;
-            }
-            catch (Exception) when (attempt < retryCount)
-            {
-                attempt++;
-                await Task.Delay(100 * attempt, cancellationToken);
-            }
-        }
+        return new BulkUpsertImporter(this, tablePath, columns, types, maxBytes);
     }
 
+
     protected override YdbTransaction BeginDbTransaction(IsolationLevel isolationLevel)
     {
         ThrowIfConnectionClosed();
diff --git a/src/Ydb.Sdk/src/Ado/YdbDataSource.cs b/src/Ydb.Sdk/src/Ado/YdbDataSource.cs
@@ -65,16 +65,6 @@ protected override async ValueTask DisposeAsyncCore() =>
         await PoolManager.ClearPool(_ydbConnectionStringBuilder.ConnectionString);
 
     protected override void Dispose(bool disposing) => DisposeAsyncCore().AsTask().GetAwaiter().GetResult();
-
-    public async Task BulkUpsertAsync(
-        string tablePath,
-        IReadOnlyList<string> columns,
-        IReadOnlyList<IReadOnlyList<object?>> rows,
-        CancellationToken cancellationToken = default)
-    {
-        await using var conn = await OpenConnectionAsync(cancellationToken);
-        await conn.BulkUpsertAsync(tablePath, columns, rows, cancellationToken);
-    }
 }
 
 #endif
diff --git a/src/Ydb.Sdk/test/Ydb.Sdk.Ado.Tests/YdbConnectionTests.cs b/src/Ydb.Sdk/test/Ydb.Sdk.Ado.Tests/YdbConnectionTests.cs
diff --git a/src/Ydb.Sdk/test/Ydb.Sdk.Ado.Tests/YdbDataSourceTests.cs b/src/Ydb.Sdk/test/Ydb.Sdk.Ado.Tests/YdbDataSourceTests.cs
