From 9512a3d9f267ec70be25e51a3a5987b6d579e252 Mon Sep 17 00:00:00 2001
From: ylorph <ylorph@users.noreply.github.com>
Date: Tue, 1 Oct 2024 16:53:10 +0200
Subject: [PATCH] Added support for  DbConnection.StateChange event handler

---
 DuckDB.NET.Data/DuckDBConnection.cs      |  5 +++
 DuckDB.NET.Test/DuckDBConnectionTests.cs | 42 ++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/DuckDB.NET.Data/DuckDBConnection.cs b/DuckDB.NET.Data/DuckDBConnection.cs
index 3c4ed8a0..fdf024ad 100644
--- a/DuckDB.NET.Data/DuckDBConnection.cs
+++ b/DuckDB.NET.Data/DuckDBConnection.cs
@@ -17,6 +17,9 @@ public partial class DuckDBConnection : DbConnection
     private DuckDBConnectionString? parsedConnection;
     private ConnectionReference? connectionReference;
     private bool inMemoryDuplication = false;
+    
+    private static readonly StateChangeEventArgs FromClosedToOpenEventArgs = new(ConnectionState.Closed, ConnectionState.Open);
+    private static readonly StateChangeEventArgs FromOpenToClosedEventArgs = new(ConnectionState.Open, ConnectionState.Closed);
 
     #region Protected Properties
 
@@ -107,6 +110,7 @@ public override void Open()
                                                   : connectionManager.GetConnectionReference(ParsedConnection);
 
         connectionState = ConnectionState.Open;
+        OnStateChange(FromClosedToOpenEventArgs);
     }
 
     protected override DbTransaction BeginDbTransaction(IsolationLevel isolationLevel)
@@ -185,6 +189,7 @@ protected override void Dispose(bool disposing)
                     connectionManager.ReturnConnectionReference(connectionReference);
                 }
                 connectionState = ConnectionState.Closed;
+                OnStateChange(FromOpenToClosedEventArgs);
             }
         }
 
diff --git a/DuckDB.NET.Test/DuckDBConnectionTests.cs b/DuckDB.NET.Test/DuckDBConnectionTests.cs
index 8f19b12e..6c48f159 100644
--- a/DuckDB.NET.Test/DuckDBConnectionTests.cs
+++ b/DuckDB.NET.Test/DuckDBConnectionTests.cs
@@ -494,4 +494,46 @@ public void ConnectionStringSetThreadsAndAccessModeOption(string accessMode, int
             value.Should().Be(threads);
         }
     }
+
+    [Fact]
+    public void ConnectionStateHandlerIsCalledOnOpen()
+    {
+        using var dbInfo = DisposableFile.GenerateInTemp("db");
+        using var connection = new DuckDBConnection(dbInfo.ConnectionString);
+        var handlerCalled = false;
+        connection.StateChange += Assert;
+        connection.Open();
+        connection.StateChange -= Assert; // otherwise the dispose/close will trigger the assert
+        
+        handlerCalled.Should().BeTrue();
+        return;
+
+        void Assert(object sender, StateChangeEventArgs args)
+        {
+            args.OriginalState.Should().Be(ConnectionState.Closed);
+            args.CurrentState.Should().Be(ConnectionState.Open);
+            handlerCalled = true;
+        }
+    }
+
+    [Fact]
+    public async Task ConnectionStateHandlerIsCalledOnClose()
+    {
+        using var dbInfo = DisposableFile.GenerateInTemp("db");
+        await using var connection = new DuckDBConnection(dbInfo.ConnectionString);
+        var handlerCalled = false;
+        await connection.OpenAsync();
+        connection.StateChange += Assert; 
+        await connection.CloseAsync();
+        
+        handlerCalled.Should().BeTrue();
+        return;
+
+        void Assert(object sender, StateChangeEventArgs args)
+        {
+            args.OriginalState.Should().Be(ConnectionState.Open);
+            args.CurrentState.Should().Be(ConnectionState.Closed);
+            handlerCalled = true;
+        }
+    }
 }
\ No newline at end of file