Add: BulkUpsertImporter

src/Ydb.Sdk/CHANGELOG.md

@@ -1,3 +1,4 @@
+- ADO.NET: Added `BulkUpsertAsync` for batch upsert operations with transaction checks and integration tests.
 - Optimization: On BadSession, do not invoke the `DeleteSession()` method.
 - Canceling AttachStream after calling the `DeleteSession` method.
 - Fixed bug: fixed issue where session was not deleted (`ClientTransportTimeout`).

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

@@ -0,0 +1,7 @@
+namespace Ydb.Sdk.Ado.BulkUpsert;
+
+public enum BulkUpsertFormat
+{
+    Proto = 0
+    // Arrow = 1
+}

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

@@ -0,0 +1,7 @@
+namespace Ydb.Sdk.Ado.BulkUpsert;
+
+public interface IBulkUpsertImporter : IAsyncDisposable
+{
+    ValueTask AddRowAsync(params Ydb.Value[] values);
+    ValueTask FlushAsync(CancellationToken cancellationToken = default);
+}

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

@@ -0,0 +1,116 @@
+using Ydb.Sdk.Value;
+
+namespace Ydb.Sdk.Ado.BulkUpsert;
+
+public sealed class YdbBulkUpsertProtoImporter : IAsyncDisposable
+{
+    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(
+        YdbConnection connection,
+        string tablePath,
+        IReadOnlyList<string> columns,
+        IReadOnlyList<Type> types,
+        int maxBytes = 1024 * 1024)
+    {
+        _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);
+        _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));
+
+        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);
+
+        var totalSize = _rows.Sum(r => r.CalculateSize());
+
+        if (totalSize >= _maxBytes)
+            await FlushAsync();
+    }
+
+    public async ValueTask AddRowsAsync(IEnumerable<YdbValue[]> rows, CancellationToken cancellationToken = default)
+    {
+        ThrowIfDisposed();
+
+        foreach (var values in rows)
+        {
+            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);
+
+            var totalSize = _rows.Sum(r => r.CalculateSize());
+
+            if (totalSize >= _maxBytes)
+                await FlushAsync(cancellationToken);
+        }
+    }
+
+    public async ValueTask DisposeAsync()
+    {
+        if (_disposed) return;
+        await FlushAsync();
+        _disposed = true;
+    }
+
+    public async ValueTask FlushAsync(CancellationToken cancellationToken = default)
+    {
+        ThrowIfDisposed();
+
+        if (_rows.Count == 0)
+            return;
+
+        await _connection.BulkUpsertProtoAsync(
+            _tablePath,
+            GetStructType(),
+            new List<Ydb.Value>(_rows),
+            cancellationToken);
+
+        _rows.Clear();
+    }
+
+    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]
+            });
+        }
+
+        return structType;
+    }
+
+    private void ThrowIfDisposed()
+    {
+        if (_disposed)
+            throw new ObjectDisposedException(nameof(YdbBulkUpsertProtoImporter));
+    }
+
+    public IReadOnlyList<Ydb.Value> GetBufferedRows() => _rows;
+}

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

@@ -1,8 +1,12 @@
 using System.Data;
 using System.Data.Common;
 using System.Diagnostics.CodeAnalysis;
+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;
 
 namespace Ydb.Sdk.Ado;
@@ -52,6 +56,92 @@ public YdbConnection(YdbConnectionStringBuilder connectionStringBuilder)
         ConnectionStringBuilder = connectionStringBuilder;
     }
 
+    public async Task BulkUpsertAsync(
+        string tablePath,
+        IReadOnlyList<string> columns,
+        IReadOnlyList<IReadOnlyList<object?>> rows,
+        CancellationToken cancellationToken = default)
+    {
+        if (CurrentTransaction is { Completed: false })
+            throw new InvalidOperationException("BulkUpsert 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 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 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 async Task BulkUpsertProtoAsync(
+        string tablePath,
+        Type structType,
+        IReadOnlyList<Ydb.Value> chunk,
+        CancellationToken cancellationToken = default,
+        int retryCount = 3)
+    {
+        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 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);
+            }
+        }
+    }
+
     protected override YdbTransaction BeginDbTransaction(IsolationLevel isolationLevel)
     {
         ThrowIfConnectionClosed();

src/Ydb.Sdk/src/Ado/YdbDataSource.cs

@@ -65,6 +65,16 @@ 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