From 18ff1496e69506db879bf24b823613fe8a6a6353 Mon Sep 17 00:00:00 2001
From: KirillKurdyukov <kurdyukov-kir@ydb.tech>
Date: Wed, 7 Aug 2024 16:17:15 +0300
Subject: [PATCH] feat: impl DbDataSource

---
 src/Ydb.Sdk/src/Ado/YdbDataSource.cs        | 69 +++++++++++++++++++++
 src/Ydb.Sdk/tests/Ado/YdbDataSourceTests.cs | 55 ++++++++++++++++
 2 files changed, 124 insertions(+)
 create mode 100644 src/Ydb.Sdk/src/Ado/YdbDataSource.cs
 create mode 100644 src/Ydb.Sdk/tests/Ado/YdbDataSourceTests.cs

diff --git a/src/Ydb.Sdk/src/Ado/YdbDataSource.cs b/src/Ydb.Sdk/src/Ado/YdbDataSource.cs
new file mode 100644
index 00000000..74e30855
--- /dev/null
+++ b/src/Ydb.Sdk/src/Ado/YdbDataSource.cs
@@ -0,0 +1,69 @@
+#if NET7_0_OR_GREATER
+
+using System.Data.Common;
+
+namespace Ydb.Sdk.Ado;
+
+public class YdbDataSource : DbDataSource
+{
+    private readonly YdbConnectionStringBuilder _ydbConnectionStringBuilder;
+
+    public YdbDataSource(YdbConnectionStringBuilder connectionStringBuilder)
+    {
+        _ydbConnectionStringBuilder = connectionStringBuilder;
+    }
+
+    public YdbDataSource(string connectionString)
+    {
+        _ydbConnectionStringBuilder = new YdbConnectionStringBuilder(connectionString);
+    }
+
+    public YdbDataSource()
+    {
+        _ydbConnectionStringBuilder = new YdbConnectionStringBuilder();
+    }
+
+    protected override YdbConnection CreateDbConnection()
+    {
+        return new YdbConnection(_ydbConnectionStringBuilder);
+    }
+
+    protected override YdbConnection OpenDbConnection()
+    {
+        var dbConnection = CreateDbConnection();
+        try
+        {
+            dbConnection.Open();
+            return dbConnection;
+        }
+        catch
+        {
+            dbConnection.Close();
+            throw;
+        }
+    }
+
+    public new YdbConnection CreateConnection() => CreateDbConnection();
+
+    public new YdbConnection OpenConnection() => OpenDbConnection();
+
+    public new async ValueTask<YdbConnection> OpenConnectionAsync(CancellationToken cancellationToken = default)
+    {
+        var ydbConnection = CreateDbConnection();
+
+        try
+        {
+            await ydbConnection.OpenAsync(cancellationToken);
+            return ydbConnection;
+        }
+        catch
+        {
+            await ydbConnection.CloseAsync();
+            throw;
+        }
+    }
+
+    public override string ConnectionString => _ydbConnectionStringBuilder.ConnectionString;
+}
+
+#endif
diff --git a/src/Ydb.Sdk/tests/Ado/YdbDataSourceTests.cs b/src/Ydb.Sdk/tests/Ado/YdbDataSourceTests.cs
new file mode 100644
index 00000000..bc4f8ebc
--- /dev/null
+++ b/src/Ydb.Sdk/tests/Ado/YdbDataSourceTests.cs
@@ -0,0 +1,55 @@
+using Xunit;
+using Ydb.Sdk.Ado;
+
+#if NET7_0_OR_GREATER
+namespace Ydb.Sdk.Tests.Ado;
+
+[Trait("Category", "Integration")]
+public class YdbDataSourceTests
+{
+    private const int SelectedCount = 100;
+
+    private readonly YdbDataSource _dataSource = new("MaxSessionPool=10");
+
+    [Fact]
+    public async Task OpenConnectionAsync_WhenMaxSessionPool10_ReturnOpenConnection()
+    {
+        var tasks = new Task[SelectedCount];
+        for (var i = 0; i < SelectedCount; i++)
+        {
+            tasks[i] = Task.Run(async () =>
+            {
+                await using var ydbConnection = await _dataSource.OpenConnectionAsync();
+                var ydbCommand = new YdbCommand(ydbConnection) { CommandText = "SELECT 1" };
+                Assert.Equal(1, await ydbCommand.ExecuteScalarAsync());
+            });
+        }
+
+        await Task.WhenAll(tasks);
+    }
+
+    [Fact]
+    public void CreateCommand_FromDataSource_ReturnDbCommand()
+    {
+        for (var i = 0; i < SelectedCount; i++)
+        {
+            Assert.Equal(1, _dataSource.CreateCommand("SELECT 1;").ExecuteScalar());
+        }
+    }
+
+    [Fact]
+    public void CreateConnection_FromDataSource_ReturnNotOpenConnection()
+    {
+        using var ydbConnection = _dataSource.CreateConnection();
+        ydbConnection.Open();
+        Assert.Equal(1, new YdbCommand(ydbConnection) { CommandText = "SELECT 1" }.ExecuteScalar());
+    }
+
+    [Fact]
+    public void OpenConnection_FromDataSource_ReturnOpenConnection()
+    {
+        using var ydbConnection = _dataSource.OpenConnection();
+        Assert.Equal(1, new YdbCommand(ydbConnection) { CommandText = "SELECT 1" }.ExecuteScalar());
+    }
+}
+#endif