bulk operations

Author: KSemenenko

Directory.Build.props

@@ -25,8 +25,8 @@
         <RepositoryUrl>https://github.com/managedcode/graphrag</RepositoryUrl>
         <PackageProjectUrl>https://github.com/managedcode/graphrag</PackageProjectUrl>
         <Product>Managed Code GraphRag</Product>
-        <Version>10.0.1</Version>
-        <PackageVersion>10.0.1</PackageVersion>
+        <Version>10.0.2</Version>
+        <PackageVersion>10.0.2</PackageVersion>
 
     </PropertyGroup>
     <PropertyGroup Condition="'$(GITHUB_ACTIONS)' == 'true'">

README.md

@@ -236,12 +236,41 @@ GraphRAG ships with a first-class Apache AGE adapter (`ManagedCode.GraphRag.Post
      }
    }
    ```
-4. **Register through DI.** `services.AddPostgresGraphStore("postgres", configure: ...)` wires up `IAgeConnectionManager`, `IAgeClientFactory`, and `IGraphStore` automatically. Pool sizing behaves like EF Core—tune it via connection-string keywords such as `Maximum Pool Size` when needed, otherwise the standard Npgsql defaults apply. The first Postgres store you register becomes the default `IGraphStore`, and additional stores stay keyed-only:
+4. **Register through DI.** `services.AddPostgresGraphStore("postgres", configure: ...)` wires up `IAgeConnectionManager`, `IAgeClientFactory`, `IGraphStore`, `IScopedGraphStore`, and `IBulkGraphStore`. Pool sizing follows the standard Npgsql settings (configure `Max Pool Size`, `Timeout`, etc. inside the connection string). The first registration becomes the default unkeyed `IGraphStore`; additional stores remain keyed-only.
+
    ```csharp
-   await using var client = ageClientFactory.CreateClient();
-   await client.OpenConnectionAsync(cancellationToken);
+   var services = new ServiceCollection()
+       .AddLogging()
+       .AddPostgresGraphStore("postgres", options =>
+       {
+           options.ConnectionString = postgresConnectionString;
+           options.GraphName = "graphrag";
+       });
+
+   await using var provider = services.BuildServiceProvider();
+
+   // Regular graph operations
+   var graphStore = provider.GetRequiredService<IGraphStore>();
+
+   // Scoped operations reuse a single AGE/Postgres connection for the lifetime of the scope
+   var scopedStore = provider.GetRequiredService<IScopedGraphStore>();
+   await using (await scopedStore.CreateScopeAsync())
+   {
+       await graphStore.UpsertNodeAsync("node-1", "Example", new Dictionary<string, object?> { ["name"] = "Scoped" });
+       await graphStore.UpsertNodeAsync("node-2", "Example", new Dictionary<string, object?> { ["name"] = "Connection" });
+   }
+
+   // Bulk helpers batch large workloads while keeping the scoped connection alive
+   var bulkStore = provider.GetRequiredService<IBulkGraphStore>();
+   await bulkStore.UpsertNodesAsync(new[]
+   {
+       new GraphNodeUpsert("bulk-1", "Example", new Dictionary<string, object?> { ["name"] = "Bulk" }),
+       new GraphNodeUpsert("bulk-2", "Example", new Dictionary<string, object?> { ["name"] = "Write" })
+   });
    ```
 
+   The `AgeConnectionManager` automatically retries transient `53300: too many clients` errors (up to three exponential backoff attempts) so scopes can wait for a free slot before failing. When a scope is disposed, the underlying `IAgeClientScope` returns its connection to the pool, keeping concurrency predictable even under heavy fan-out.
+
 ### Neo4j Setup
 
 Neo4j support lives in `ManagedCode.GraphRag.Neo4j` and uses the official Bolt driver:

src/ManagedCode.GraphRag.Postgres/ApacheAge/AGEClientEventId.cs

@@ -5,6 +5,7 @@ internal static class AgeClientEventId
     #region Connection
     public const int CONNECTION_OPENED = 1000;
     public const int CONNECTION_CLOSED = 1001;
+    public const int CONNECTION_RETRYING = 1002;
 
     public const int NULL_CONNECTION_WARNING = 1800;
 

src/ManagedCode.GraphRag.Postgres/ApacheAge/AgeClientFactory.cs

@@ -1,21 +1,152 @@
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
+using Npgsql;
 
 namespace GraphRag.Storage.Postgres.ApacheAge;
 
 public interface IAgeClientFactory
 {
-    AgeClient CreateClient();
+    IAgeClient CreateClient();
+    ValueTask<IAgeClientScope> CreateScopeAsync(CancellationToken cancellationToken = default);
 }
 
 internal sealed class AgeClientFactory([FromKeyedServices] IAgeConnectionManager connectionManager, ILoggerFactory loggerFactory) : IAgeClientFactory
 {
     private readonly IAgeConnectionManager _connectionManager = connectionManager ?? throw new ArgumentNullException(nameof(connectionManager));
     private readonly ILoggerFactory _loggerFactory = loggerFactory ?? throw new ArgumentNullException(nameof(loggerFactory));
+    private readonly AsyncLocal<AgeClientScopeState?> _currentScope = new();
 
-    public AgeClient CreateClient()
+    public IAgeClient CreateClient()
+    {
+        if (_currentScope.Value is { } scopeState)
+        {
+            return scopeState.CreateLease();
+        }
+
+        return CreatePhysicalClient();
+    }
+
+    public ValueTask<IAgeClientScope> CreateScopeAsync(CancellationToken cancellationToken = default)
+    {
+        cancellationToken.ThrowIfCancellationRequested();
+
+        if (_currentScope.Value is not null)
+        {
+            throw new InvalidOperationException("An AGE client scope is already active for this asynchronous context.");
+        }
+
+        var scopeClient = CreatePhysicalClient();
+        var scopeState = new AgeClientScopeState(scopeClient);
+        _currentScope.Value = scopeState;
+
+        return new ValueTask<IAgeClientScope>(new ActiveAgeClientScope(this, scopeState));
+    }
+
+    private AgeClient CreatePhysicalClient()
     {
         var logger = _loggerFactory.CreateLogger<AgeClient>();
         return new AgeClient(_connectionManager, logger);
     }
+
+    private void ClearScope(AgeClientScopeState state)
+    {
+        if (!ReferenceEquals(_currentScope.Value, state))
+        {
+            return;
+        }
+
+        _currentScope.Value = null;
+    }
+
+    private sealed class ActiveAgeClientScope(AgeClientFactory factory, AgeClientFactory.AgeClientScopeState state) : IAgeClientScope
+    {
+        private readonly AgeClientFactory _factory = factory;
+        private readonly AgeClientScopeState _state = state;
+        private bool _disposed;
+
+        public async ValueTask DisposeAsync()
+        {
+            if (_disposed)
+            {
+                return;
+            }
+
+            if (_state.HasActiveLease)
+            {
+                throw new InvalidOperationException("Cannot dispose an AGE client scope while an operation is still running.");
+            }
+
+            _disposed = true;
+            _factory.ClearScope(_state);
+            await _state.Client.CloseConnectionAsync().ConfigureAwait(false);
+            await _state.Client.DisposeAsync().ConfigureAwait(false);
+        }
+    }
+
+    private sealed class AgeClientScopeState(AgeClient client)
+    {
+        private int _leaseActive;
+
+        public AgeClient Client { get; } = client;
+
+        public IAgeClient CreateLease()
+        {
+            if (Interlocked.CompareExchange(ref _leaseActive, 1, 0) == 1)
+            {
+                throw new InvalidOperationException("An AGE client scope cannot be used concurrently. Await ongoing operations before issuing new ones within the same scope.");
+            }
+
+            return new ScopedAgeClient(this);
+        }
+
+        public void ReleaseLease()
+        {
+            Volatile.Write(ref _leaseActive, 0);
+        }
+
+        public bool HasActiveLease => Volatile.Read(ref _leaseActive) == 1;
+    }
+
+    private sealed class ScopedAgeClient(AgeClientFactory.AgeClientScopeState state) : IAgeClient
+    {
+        private readonly AgeClientScopeState _state = state;
+        private bool _disposed;
+
+        public bool IsConnected => _state.Client.IsConnected;
+
+        public NpgsqlConnection Connection => _state.Client.Connection;
+
+        public Task OpenConnectionAsync(CancellationToken cancellationToken = default) =>
+            _state.Client.OpenConnectionAsync(cancellationToken);
+
+        public Task CreateGraphAsync(string graphName, CancellationToken cancellationToken = default) =>
+            _state.Client.CreateGraphAsync(graphName, cancellationToken);
+
+        public Task DropGraphAsync(string graphName, bool cascade = false, CancellationToken cancellationToken = default) =>
+            _state.Client.DropGraphAsync(graphName, cascade, cancellationToken);
+
+        public Task ExecuteCypherAsync(string graph, string cypher, CancellationToken cancellationToken = default) =>
+            _state.Client.ExecuteCypherAsync(graph, cypher, cancellationToken);
+
+        public Task<AgeDataReader> ExecuteQueryAsync(string query, CancellationToken cancellationToken = default, params object?[] parameters) =>
+            _state.Client.ExecuteQueryAsync(query, cancellationToken, parameters);
+
+        public Task CloseConnectionAsync(CancellationToken cancellationToken = default) =>
+            Task.CompletedTask;
+
+        public Task<bool> GraphExistsAsync(string graphName, CancellationToken cancellationToken = default) =>
+            _state.Client.GraphExistsAsync(graphName, cancellationToken);
+
+        public ValueTask DisposeAsync()
+        {
+            if (_disposed)
+            {
+                return ValueTask.CompletedTask;
+            }
+
+            _disposed = true;
+            _state.ReleaseLease();
+            return ValueTask.CompletedTask;
+        }
+    }
 }

src/ManagedCode.GraphRag.Postgres/ApacheAge/AgeConnectionManager.cs

@@ -14,6 +14,10 @@ public interface IAgeConnectionManager : IAsyncDisposable, IDisposable
 
 public sealed class AgeConnectionManager : IAgeConnectionManager
 {
+    private const int ConnectionLimitMaxAttempts = 3;
+    private static readonly TimeSpan ConnectionLimitBaseDelay = TimeSpan.FromMilliseconds(200);
+    private static readonly TimeSpan ConnectionLimitMaxDelay = TimeSpan.FromSeconds(2);
+
     private readonly NpgsqlDataSource _dataSource;
     private readonly ILogger<AgeConnectionManager> _logger;
     private volatile bool _extensionEnsured;
@@ -50,7 +54,7 @@ public async Task<NpgsqlConnection> OpenConnectionAsync(CancellationToken cancel
         ThrowIfDisposed();
 
         await EnsureExtensionCreatedAsync(cancellationToken).ConfigureAwait(false);
-        var connection = await _dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
+        var connection = await OpenDataSourceConnectionAsync(cancellationToken).ConfigureAwait(false);
         await LoadAgeAsync(connection, cancellationToken).ConfigureAwait(false);
         await SetSearchPathAsync(connection, cancellationToken).ConfigureAwait(false);
         return connection;
@@ -104,8 +108,7 @@ private async Task EnsureExtensionCreatedAsync(CancellationToken cancellationTok
             return;
         }
 
-        await using var connection = new NpgsqlConnection(ConnectionString);
-        await connection.OpenAsync(cancellationToken).ConfigureAwait(false);
+        await using var connection = await OpenDataSourceConnectionAsync(cancellationToken).ConfigureAwait(false);
         await using var command = connection.CreateCommand();
         command.CommandText = "CREATE EXTENSION IF NOT EXISTS age;";
         await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
@@ -159,4 +162,37 @@ private async Task SetSearchPathAsync(NpgsqlConnection connection, CancellationT
 
     private void ThrowIfDisposed() =>
         ObjectDisposedException.ThrowIf(_disposed, nameof(AgeConnectionManager));
+
+    private async Task<NpgsqlConnection> OpenDataSourceConnectionAsync(CancellationToken cancellationToken)
+    {
+        var attempt = 0;
+
+        while (true)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+            attempt++;
+
+            try
+            {
+                return await _dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);
+            }
+            catch (PostgresException ex) when (ShouldRetry(ex, attempt))
+            {
+                var delay = GetRetryDelay(attempt);
+                LogMessages.ConnectionRetrying(_logger, ConnectionString, attempt, delay, ex.MessageText);
+                await Task.Delay(delay, cancellationToken).ConfigureAwait(false);
+            }
+        }
+    }
+
+    private static bool ShouldRetry(PostgresException ex, int attempt) =>
+        ex.SqlState == PostgresErrorCodes.TooManyConnections && attempt < ConnectionLimitMaxAttempts;
+
+    private static TimeSpan GetRetryDelay(int attempt)
+    {
+        var delayMillis = Math.Min(
+            ConnectionLimitBaseDelay.TotalMilliseconds * Math.Pow(2, attempt - 1),
+            ConnectionLimitMaxDelay.TotalMilliseconds);
+        return TimeSpan.FromMilliseconds(delayMillis);
+    }
 }

src/ManagedCode.GraphRag.Postgres/ApacheAge/IAgeClientScope.cs

@@ -0,0 +1,8 @@
+namespace GraphRag.Storage.Postgres.ApacheAge;
+
+/// <summary>
+/// Represents a scope that keeps an AGE connection alive for the lifetime of the scope.
+/// </summary>
+public interface IAgeClientScope : IAsyncDisposable
+{
+}

src/ManagedCode.GraphRag.Postgres/ApacheAge/LogMessages.cs

@@ -22,6 +22,17 @@ public static partial void ConnectionClosed(
         ILogger logger,
         string connectionString);
 
+    [LoggerMessage(
+    EventId = AgeClientEventId.CONNECTION_RETRYING,
+    Level = LogLevel.Warning,
+    Message = "Connection limit reached for {connectionString}. Attempt {attempt} failed; retrying after {delay}. Reason: {reason}")]
+    public static partial void ConnectionRetrying(
+        ILogger logger,
+        string connectionString,
+        int attempt,
+        TimeSpan delay,
+        string reason);
+
     [LoggerMessage(
     EventId = AgeClientEventId.NULL_CONNECTION_WARNING,
     Level = LogLevel.Warning,