Add Apache AGE support with new types and integration methods

Author: KSemenenko

AGENTS.md

@@ -8,6 +8,7 @@
 - Register language models through Microsoft.Extensions.AI keyed services; avoid bespoke `LanguageModelConfig` providers.
 - Always run `dotnet format GraphRag.slnx` before finishing work.
 - Always run `dotnet test GraphRag.slnx` before finishing work, after building.
+- Avoid per-client connection locks (e.g., `_connectionLock` in `AgeClient`); rely on the smart connection manager for concurrency.
 
 # Conversations
 any resulting updates to agents.md should go under the section "## Rules to follow"

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>0.0.4</Version>
-        <PackageVersion>0.0.4</PackageVersion>
+        <Version>0.0.5</Version>
+        <PackageVersion>0.0.5</PackageVersion>
 
     </PropertyGroup>
     <PropertyGroup Condition="'$(GITHUB_ACTIONS)' == 'true'">

README.md

@@ -205,6 +205,49 @@ dotnet test tests/ManagedCode.GraphRag.Tests/ManagedCode.GraphRag.Tests.csproj \
 
 ---
 
+## Apache AGE / PostgreSQL Setup
+
+GraphRAG ships with a first-class Apache AGE adapter (`ManagedCode.GraphRag.Postgres`). AGE is enabled on top of PostgreSQL, so you only need a standard Postgres instance with the AGE extension installed.
+
+1. **Run an AGE-enabled Postgres instance.** The integration tests use the official container and you can do the same locally:
+   ```bash
+   docker run --rm \
+     -e POSTGRES_USER=postgres \
+     -e POSTGRES_PASSWORD=postgres \
+     -e POSTGRES_DB=graphrag \
+     -p 5432:5432 \
+     apache/age:latest
+   ```
+2. **Provide a connection string.** `AgeConnectionManager` accepts a standard Npgsql-style string (for example `Host=localhost;Port=5432;Username=postgres;Password=postgres;Database=graphrag`). The manager automatically runs `CREATE EXTENSION IF NOT EXISTS age;`, `LOAD 'age';`, and `SET search_path = ag_catalog, "$user", public;` before any query executes.
+3. **Configure the store.** Either bind `PostgresGraphStoreOptions` in code or use configuration. The snippet below shows the JSON shape (environment variables can follow the same hierarchy, e.g. `GraphRag__GraphStores__postgres__ConnectionString`):
+   ```json
+   {
+     "GraphRag": {
+       "GraphStores": {
+         "postgres": {
+           "ConnectionString": "Host=localhost;Port=5432;Username=postgres;Password=postgres;Database=graphrag",
+           "GraphName": "graphrag",
+           "MaxConnections": 40,
+           "MakeDefault": true
+         }
+       }
+     }
+   }
+   ```
+4. **Register through DI.** `services.AddPostgresGraphStore("postgres", configure: ...)` wires up `IAgeConnectionManager`, `IAgeClientFactory`, and `IGraphStore` automatically. `MaxConnections` caps the number of concurrent AGE sessions (the default is 40 so you stay under the container’s `max_connections`).
+
+> **Tip:** `IGraphStore` now exposes `GetNodesAsync` and `GetRelationshipsAsync` in addition to the targeted APIs (`InitializeAsync`, `Upsert*`, `GetOutgoingRelationshipsAsync`). These use the new AGE-powered enumerations so you can inspect or export the full graph without dropping down to concrete implementations.
+
+> **Pagination:** `GetNodesAsync` and `GetRelationshipsAsync` accept an optional `GraphTraversalOptions` object (`new GraphTraversalOptions { Skip = 100, Take = 50 }`) if you want to page through very large graphs. The defaults stream everything, one record at a time, without materialising the entire graph in memory.
+
+---
+
+## Credits
+
+- **pg-age** ([Allison-E/pg-age](https://github.com/Allison-E/pg-age)) — we vendor this Apache AGE client library (see `src/ManagedCode.GraphRag.Postgres/ApacheAge`) so GraphRAG for .NET can rely on a battle-tested connector. Many thanks to Allison and contributors for making AGE on PostgreSQL accessible.
+
+---
+
 ## Additional Documentation & Diagrams
 
 - [`docs/indexing-and-query.md`](docs/indexing-and-query.md) explains how each workflow maps to the GraphRAG research diagrams (default data flow, query orchestrations, prompt tuning strategies) published at [microsoft.github.io/graphrag](https://microsoft.github.io/graphrag/).

src/ManagedCode.GraphRag.CosmosDb/CosmosGraphStore.cs

@@ -69,4 +69,72 @@ async IAsyncEnumerable<GraphRelationship> Fetch([EnumeratorCancellation] Cancell
     private sealed record NodeDocument(string Id, string Label, Dictionary<string, object?> Properties);
 
     private sealed record EdgeDocument(string Id, string SourceId, string TargetId, string Type, Dictionary<string, object?> Properties);
+
+    public IAsyncEnumerable<GraphNode> GetNodesAsync(GraphTraversalOptions? options = null, CancellationToken cancellationToken = default)
+    {
+        options?.Validate();
+        return Fetch(options, cancellationToken);
+
+        async IAsyncEnumerable<GraphNode> Fetch(GraphTraversalOptions? traversalOptions, [EnumeratorCancellation] CancellationToken token = default)
+        {
+            var container = _client.GetContainer(_databaseId, _nodesContainerId);
+            IQueryable<NodeDocument> queryable = container.GetItemLinqQueryable<NodeDocument>(allowSynchronousQueryExecution: false);
+            queryable = queryable.OrderBy(node => node.Id);
+
+            if (traversalOptions?.Skip is > 0 and var skip)
+            {
+                queryable = queryable.Skip(skip);
+            }
+
+            if (traversalOptions?.Take is { } take)
+            {
+                queryable = queryable.Take(take);
+            }
+
+            using var iterator = queryable.ToFeedIterator();
+
+            while (iterator.HasMoreResults)
+            {
+                var response = await iterator.ReadNextAsync(token);
+                foreach (var node in response)
+                {
+                    yield return new GraphNode(node.Id, node.Label, node.Properties);
+                }
+            }
+        }
+    }
+
+    public IAsyncEnumerable<GraphRelationship> GetRelationshipsAsync(GraphTraversalOptions? options = null, CancellationToken cancellationToken = default)
+    {
+        options?.Validate();
+        return FetchEdges(options, cancellationToken);
+
+        async IAsyncEnumerable<GraphRelationship> FetchEdges(GraphTraversalOptions? traversalOptions, [EnumeratorCancellation] CancellationToken token = default)
+        {
+            var container = _client.GetContainer(_databaseId, _edgesContainerId);
+            IQueryable<EdgeDocument> queryable = container.GetItemLinqQueryable<EdgeDocument>(allowSynchronousQueryExecution: false);
+            queryable = queryable.OrderBy(edge => edge.SourceId).ThenBy(edge => edge.TargetId);
+
+            if (traversalOptions?.Skip is > 0 and var skip)
+            {
+                queryable = queryable.Skip(skip);
+            }
+
+            if (traversalOptions?.Take is { } take)
+            {
+                queryable = queryable.Take(take);
+            }
+
+            using var iterator = queryable.ToFeedIterator();
+
+            while (iterator.HasMoreResults)
+            {
+                var response = await iterator.ReadNextAsync(token);
+                foreach (var edge in response)
+                {
+                    yield return new GraphRelationship(edge.SourceId, edge.TargetId, edge.Type, edge.Properties);
+                }
+            }
+        }
+    }
 }

src/ManagedCode.GraphRag.Neo4j/Neo4jGraphStore.cs

@@ -77,6 +77,88 @@ async IAsyncEnumerable<GraphRelationship> Fetch([EnumeratorCancellation] Cancell
         }
     }
 
+    public IAsyncEnumerable<GraphRelationship> GetRelationshipsAsync(GraphTraversalOptions? options = null, CancellationToken cancellationToken = default)
+    {
+        options?.Validate();
+        cancellationToken.ThrowIfCancellationRequested();
+        return FetchAll(options, cancellationToken);
+
+        async IAsyncEnumerable<GraphRelationship> FetchAll(GraphTraversalOptions? traversalOptions, [EnumeratorCancellation] CancellationToken token = default)
+        {
+            await using var session = _driver.AsyncSession();
+            var parameters = new Dictionary<string, object>();
+            var pagination = BuildNeo4jPaginationClause(traversalOptions, parameters, orderBy: "source.id, target.id, type(rel)");
+            var cursor = await session.RunAsync(
+                @"MATCH (source)-[rel]->(target)
+                  RETURN source.id AS SourceId, target.id AS TargetId, type(rel) AS Type, properties(rel) AS Properties" + pagination,
+                parameters);
+
+            while (await cursor.FetchAsync())
+            {
+                token.ThrowIfCancellationRequested();
+                var record = cursor.Current;
+                var properties = record["Properties"].As<IDictionary<string, object?>>();
+                yield return new GraphRelationship(
+                    record["SourceId"].As<string>(),
+                    record["TargetId"].As<string>(),
+                    record["Type"].As<string>(),
+                    properties.ToDictionary(kvp => kvp.Key, kvp => kvp.Value));
+            }
+        }
+    }
+
+    public IAsyncEnumerable<GraphNode> GetNodesAsync(GraphTraversalOptions? options = null, CancellationToken cancellationToken = default)
+    {
+        options?.Validate();
+        cancellationToken.ThrowIfCancellationRequested();
+        return FetchNodes(options, cancellationToken);
+
+        async IAsyncEnumerable<GraphNode> FetchNodes(GraphTraversalOptions? traversalOptions, [EnumeratorCancellation] CancellationToken token = default)
+        {
+            await using var session = _driver.AsyncSession();
+            var parameters = new Dictionary<string, object>();
+            var pagination = BuildNeo4jPaginationClause(traversalOptions, parameters, orderBy: "n.id");
+            var cursor = await session.RunAsync(
+                @"MATCH (n)
+                  RETURN head(labels(n)) AS Label, n.id AS Id, properties(n) AS Properties" + pagination,
+                parameters);
+
+            while (await cursor.FetchAsync())
+            {
+                token.ThrowIfCancellationRequested();
+                var record = cursor.Current;
+                var properties = record["Properties"].As<IDictionary<string, object?>>();
+                yield return new GraphNode(
+                    record["Id"].As<string>(),
+                    record["Label"].As<string>(),
+                    properties.ToDictionary(kvp => kvp.Key, kvp => kvp.Value));
+            }
+        }
+    }
+
+    private static string BuildNeo4jPaginationClause(GraphTraversalOptions? options, IDictionary<string, object> parameters, string orderBy)
+    {
+        var clause = $" ORDER BY {orderBy}";
+        if (options is null)
+        {
+            return clause;
+        }
+
+        if (options.Skip is > 0 and var skip)
+        {
+            parameters["skip"] = skip;
+            clause += " SKIP $skip";
+        }
+
+        if (options.Take is { } take)
+        {
+            parameters["limit"] = take;
+            clause += " LIMIT $limit";
+        }
+
+        return clause;
+    }
+
     private static string EscapeLabel(string value)
     {
         if (value.Any(ch => !char.IsLetterOrDigit(ch) && ch != '_' && ch != ':'))

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

@@ -0,0 +1,42 @@
+namespace GraphRag.Storage.Postgres.ApacheAge;
+
+internal static class AgeClientEventId
+{
+    #region Connection
+    public const int CONNECTION_OPENED = 1000;
+    public const int CONNECTION_CLOSED = 1001;
+
+    public const int NULL_CONNECTION_WARNING = 1800;
+
+    public const int OPEN_CONNECTION_ERROR = 1900;
+    public const int CLOSE_CONNECTION_ERROR = 1901;
+    #endregion
+
+    #region Internals
+    public const int EXTENSION_CREATED = 2101;
+    public const int EXTENSION_LOADED = 2103;
+    public const int RETRIEVED_CURRENT_SEARCH_PATH = 2104;
+    public const int AG_CATALOG_ADDED_TO_SEARCH_PATH = 2105;
+
+    public const int EXTENSION_NOT_CREATED_ERROR = 2900;
+    public const int EXTENSION_NOT_LOADED_ERROR = 2901;
+    public const int AG_CATALOG_NOT_ADDED_TO_SEARCH_PATH_ERROR = 2902;
+    #endregion
+
+    #region Commands
+    public const int GRAPH_CREATED = 3001;
+    public const int GRAPH_DROPPED = 3002;
+    public const int GRAPH_EXISTS = 3003;
+    public const int GRAPH_DOES_NOT_EXIST = 3004;
+
+    public const int CYPHER_EXECUTED = 3101;
+    public const int QUERY_EXECUTED = 3102;
+
+    public const int GRAPH_NOT_CREATED_ERROR = 3900;
+    public const int GRAPH_NOT_DROPPED_ERROR = 3901;
+    public const int CYPHER_EXECUTION_ERROR = 3902;
+    public const int QUERY_EXECUTION_ERROR = 3903;
+    #endregion
+
+    public const int UNKNOWN_ERROR = 9900;
+}