Add: BulkUpsertImporter

src/Ydb.Sdk/CHANGELOG.md

@@ -1,3 +1,4 @@
+- ADO.NET: Added `BeginBulkUpsertImport` 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/BulkUpsertImporter.cs

@@ -0,0 +1,98 @@
+using Google.Protobuf.Collections;
+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 IDriver _driver;
+    private readonly string _tablePath;
+    private readonly IReadOnlyList<string> _columns;
+    private readonly int _maxBytes;
+    private readonly RepeatedField<Ydb.Value> _rows = new();
+    private readonly CancellationToken _cancellationToken;
+    private StructType? _structType;
+    private int _currentBytes;
+
+    public BulkUpsertImporter(
+        IDriver driver,
+        string tableName,
+        IReadOnlyList<string> columns,
+        CancellationToken cancellationToken = default,
+        int maxBytes = 64 * 1024 * 1024)
+    {
+        _driver = driver;
+        _tablePath = tableName;
+        _columns = columns;
+        _maxBytes = maxBytes;
+        _cancellationToken = cancellationToken;
+        _maxBytes = maxBytes / 2;
+    }
+
+    public async ValueTask AddRowAsync(object?[] values)
+    {
+        if (values.Length != _columns.Count)
+            throw new ArgumentException("Values count must match columns count", nameof(values));
+
+        var ydbValues = values.Select(v =>
+            v as YdbValue ?? (v is YdbParameter param ? param.YdbValue : new YdbParameter { Value = v }.YdbValue)
+        ).ToArray();
+
+        var protoStruct = new Ydb.Value();
+        foreach (var value in ydbValues)
+            protoStruct.Items.Add(value.GetProto().Value);
+
+        var rowSize = protoStruct.CalculateSize();
+
+        if (_currentBytes + rowSize > _maxBytes && _rows.Count > 0)
+        {
+            await FlushAsync(_cancellationToken);
+        }
+
+        _rows.Add(protoStruct);
+        _currentBytes += rowSize;
+
+        _structType ??= new StructType
+        {
+            Members =
+            {
+                _columns.Select((col, i) => new StructMember
+                {
+                    Name = col,
+                    Type = ydbValues[i].GetProto().Type
+                })
+            }
+        };
+    }
+
+    public async ValueTask FlushAsync(CancellationToken cancellationToken = default)
+    {
+        if (_rows.Count == 0) return;
+        if (_structType == null)
+            throw new InvalidOperationException("structType is undefined");
+
+        var typedValue = new TypedValue
+        {
+            Type = new Type { ListType = new ListType { Item = new Type { StructType = _structType } } },
+            Value = new Ydb.Value { Items = { _rows } }
+        };
+
+        var req = new BulkUpsertRequest { Table = _tablePath, Rows = typedValue };
+
+        var resp = await _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();
+        _currentBytes = 0;
+    }
+}

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

@@ -0,0 +1,7 @@
+namespace Ydb.Sdk.Ado.BulkUpsert;
+
+public interface IBulkUpsertImporter
+{
+    ValueTask AddRowAsync(object?[] rows);
+    ValueTask FlushAsync(CancellationToken cancellationToken = default);
+}

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

@@ -1,6 +1,7 @@
 using System.Data;
 using System.Data.Common;
 using System.Diagnostics.CodeAnalysis;
+using Ydb.Sdk.Ado.BulkUpsert;
 using Ydb.Sdk.Ado.Session;
 using Ydb.Sdk.Services.Query;
 using static System.Data.IsolationLevel;
@@ -52,6 +53,22 @@ public YdbConnection(YdbConnectionStringBuilder connectionStringBuilder)
         ConnectionStringBuilder = connectionStringBuilder;
     }
 
+    public IBulkUpsertImporter BeginBulkUpsertImport(
+        string name,
+        IReadOnlyList<string> columns,
+        CancellationToken cancellationToken = default
+    )
+    {
+        ThrowIfConnectionClosed();
+        if (CurrentTransaction is { Completed: false })
+            throw new InvalidOperationException("BulkUpsert cannot be used inside an active transaction.");
+
+        var database = ConnectionStringBuilder.Database.TrimEnd('/');
+        var tablePath = string.IsNullOrEmpty(database) ? name : $"{database}/{name}";
+
+        return new BulkUpsertImporter(Session.Driver, tablePath, columns, cancellationToken);
+    }
+
     protected override YdbTransaction BeginDbTransaction(IsolationLevel isolationLevel)
     {
         ThrowIfConnectionClosed();

src/Ydb.Sdk/test/Ydb.Sdk.Ado.Tests/YdbConnectionTests.cs

@@ -308,4 +308,187 @@ private List<Task> GenerateTasks() => Enumerable.Range(0, 100).Select(async i =>
 
     protected override async Task OnDisposeAsync() =>
         await YdbConnection.ClearPool(new YdbConnection(_connectionStringTls));
+
+    [Fact]
+    public async Task BulkUpsertImporter_HappyPath_Add_Flush()
+    {
+        var tableName = $"BulkImporter_{Guid.NewGuid():N}";
+
+        var conn = new YdbConnection(_connectionStringTls);
+        await conn.OpenAsync();
+        try
+        {
+            await using (var createCmd = conn.CreateCommand())
+            {
+                createCmd.CommandText = $@"
+                    CREATE TABLE {tableName} (
+                        Id Int32,
+                        Name Utf8,
+                        PRIMARY KEY (Id)
+                    )";
+                await createCmd.ExecuteNonQueryAsync();
+            }
+
+            var columns = new[] { "Id", "Name" };
+
+            var importer = conn.BeginBulkUpsertImport(tableName, columns);
+
+            await importer.AddRowAsync([YdbValue.MakeInt32(1), YdbValue.MakeUtf8("Alice")]);
+            await importer.AddRowAsync([YdbValue.MakeInt32(2), YdbValue.MakeUtf8("Bob")]);
+            await importer.FlushAsync();
+
+            await using (var checkCmd = conn.CreateCommand())
+            {
+                checkCmd.CommandText = $"SELECT COUNT(*) FROM {tableName}";
+                var count = Convert.ToInt32(await checkCmd.ExecuteScalarAsync());
+                Assert.Equal(2, count);
+            }
+
+            importer = conn.BeginBulkUpsertImport(tableName, columns);
+            await importer.AddRowAsync([YdbValue.MakeInt32(3), YdbValue.MakeUtf8("Charlie")]);
+            await importer.AddRowAsync([YdbValue.MakeInt32(4), YdbValue.MakeUtf8("Diana")]);
+            await importer.FlushAsync();
+
+            await using (var checkCmd = conn.CreateCommand())
+            {
+                checkCmd.CommandText = $"SELECT Name FROM {tableName} ORDER BY Id";
+                var names = new List<string>();
+                await using var reader = await checkCmd.ExecuteReaderAsync();
+                while (await reader.ReadAsync())
+                    names.Add(reader.GetString(0));
+                Assert.Contains("Alice", names);
+                Assert.Contains("Bob", names);
+                Assert.Contains("Charlie", names);
+                Assert.Contains("Diana", names);
+            }
+        }
+        finally
+        {
+            await using var dropCmd = conn.CreateCommand();
+            dropCmd.CommandText = $"DROP TABLE {tableName}";
+            await dropCmd.ExecuteNonQueryAsync();
+        }
+    }
+
+    [Fact]
+    public async Task BulkUpsertImporter_ThrowsOnInvalidRowCount()
+    {
+        var tableName = $"BulkImporter_{Guid.NewGuid():N}";
+        var conn = new YdbConnection(_connectionStringTls);
+        await conn.OpenAsync();
+        try
+        {
+            await using (var createCmd = conn.CreateCommand())
+            {
+                createCmd.CommandText = $@"
+                CREATE TABLE {tableName} (
+                    Id Int32,
+                    Name Utf8,
+                    PRIMARY KEY (Id)
+                )";
+                await createCmd.ExecuteNonQueryAsync();
+            }
+
+            var columns = new[] { "Id", "Name" };
+
+            var importer = conn.BeginBulkUpsertImport(tableName, columns);
+
+            var badRow = new object?[] { YdbValue.MakeInt32(1) };
+            await Assert.ThrowsAsync<ArgumentException>(async () => await importer.AddRowAsync([badRow]));
+
+            await Assert.ThrowsAsync<ArgumentException>(async () =>
+            {
+                await importer.AddRowAsync([
+                    new object?[] { YdbValue.MakeInt32(2) }
+                ]);
+            });
+        }
+        finally
+        {
+            await using var dropCmd = conn.CreateCommand();
+            dropCmd.CommandText = $"DROP TABLE {tableName}";
+            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 conn = new YdbConnection(_connectionStringTls);
+        await conn.OpenAsync();
+        try
+        {
+            foreach (var table in new[] { table1, table2 })
+            {
+                await using var createCmd = conn.CreateCommand();
+                createCmd.CommandText = $@"CREATE TABLE {table} (
+                    Id Int32,
+                    Name Utf8,
+                    PRIMARY KEY (Id)
+                )";
+                await createCmd.ExecuteNonQueryAsync();
+            }
+
+            var columns = new[] { "Id", "Name" };
+
+            await Task.WhenAll(
+                Task.Run(async () =>
+                {
+                    var importer = conn.BeginBulkUpsertImport(table1, columns);
+                    var rows = Enumerable.Range(0, 20)
+                        .Select(i => new object?[] { YdbValue.MakeInt32(i), YdbValue.MakeUtf8($"A{i}") })
+                        .ToArray();
+                    foreach (var row in rows)
+                        await importer.AddRowAsync(row);
+                    await importer.FlushAsync();
+                }),
+                Task.Run(async () =>
+                {
+                    var importer = conn.BeginBulkUpsertImport(table2, columns);
+                    var rows = Enumerable.Range(0, 20)
+                        .Select(i => new object?[] { YdbValue.MakeInt32(i), YdbValue.MakeUtf8($"B{i}") })
+                        .ToArray();
+                    foreach (var row in rows)
+                        await importer.AddRowAsync(row);
+                    await importer.FlushAsync();
+                })
+            );
+
+            foreach (var table in new[] { table1, table2 })
+            {
+                await using var checkCmd = conn.CreateCommand();
+                checkCmd.CommandText = $"SELECT COUNT(*) FROM {table}";
+                var count = Convert.ToInt32(await checkCmd.ExecuteScalarAsync());
+                Assert.Equal(20, count);
+            }
+        }
+        finally
+        {
+            foreach (var table in new[] { table1, table2 })
+            {
+                await using var dropCmd = conn.CreateCommand();
+                dropCmd.CommandText = $"DROP TABLE {table}";
+                await dropCmd.ExecuteNonQueryAsync();
+            }
+        }
+    }
+
+    [Fact]
+    public async Task BulkUpsertImporter_ThrowsOnNonexistentTable()
+    {
+        var tableName = $"Nonexistent_{Guid.NewGuid():N}";
+        var conn = new YdbConnection(_connectionStringTls);
+        await conn.OpenAsync();
+
+        var columns = new[] { "Id", "Name" };
+
+        var importer = conn.BeginBulkUpsertImport(tableName, columns);
+
+        await importer.AddRowAsync([YdbValue.MakeInt32(1), YdbValue.MakeUtf8("NotExists")]);
+
+        await Assert.ThrowsAsync<YdbException>(async () => { await importer.FlushAsync(); });
+    }
 }