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,15 @@
+using Ydb.Sdk.Value;
+
+namespace Ydb.Sdk.Ado.BulkUpsert;
+
+public interface IBulkUpsertImporter : IAsyncDisposable
+{
+    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();
+}

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

@@ -0,0 +1,126 @@
+using Ydb.Sdk.Value;
+
+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 int _maxBytes;
+    private readonly List<Ydb.Value> _rows = new();
+    private bool _disposed;
+
+    public BulkUpsertImporter(
+        YdbConnection connection,
+        string tablePath,
+        IReadOnlyList<string> columns,
+        IReadOnlyList<Type> types,
+        int maxBytes = 1024 * 1024)
+    {
+        _connection = connection;
+        _tablePath = tablePath;
+        _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));
+
+        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 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)
+    {
+        ThrowIfDisposed();
+
+        foreach (var values in rows)
+            await AddRowAsync(values);
+    }
+
+    public async ValueTask AddRowsAsync(IEnumerable<object?[]> rows, CancellationToken cancellationToken = default)
+    {
+        ThrowIfDisposed();
+
+        foreach (var values in rows)
+            await AddRowAsync(values);
+    }
+
+    public async ValueTask FlushAsync(CancellationToken cancellationToken = default)
+    {
+        ThrowIfDisposed();
+        if (_rows.Count == 0) return;
+
+        await _connection.BulkUpsertProtoAsync(_tablePath, GetStructType(), _rows.ToList(), cancellationToken);
+        _rows.Clear();
+    }
+
+    public IReadOnlyList<Ydb.Value> GetBufferedRows() => _rows;
+
+    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] });
+        return structType;
+    }
+
+    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/YdbConnection.cs

@@ -1,8 +1,12 @@
 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.Table;
+using Ydb.Table.V1;
 using static System.Data.IsolationLevel;
 
 namespace Ydb.Sdk.Ado;
@@ -52,6 +56,46 @@ 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,
+        IReadOnlyList<string> columns,
+        IReadOnlyList<Type> types,
+        int maxBytes = 1024 * 1024)
+    {
+        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);
+    }
+
+
     protected override YdbTransaction BeginDbTransaction(IsolationLevel isolationLevel)
     {
         ThrowIfConnectionClosed();