Skip to content

Commit b89bd66

Browse files
MWG-LoganMWG-Logan
authored
Mcp/Storage fixes (MWG-Logan#38)
* Remove Azure Table Storage service registration Removed Azure Table Storage registration from services. * Refactor TableStorageService constructor for config Refactored constructor to initialize TableServiceClient based on environment variables for better configuration handling. * Refactor namespace and add Azure.Identity using * Move KnowledgeGraphService to Services namespace * Move namespace for RelationService to Services * Add using directives for Services and Storage * Add Storage namespace to RelationService * Add using directive for Storage namespace * Refactor GraphFunctions to use constructor parameters * Refactor relation ID handling in RelationService * Refactor RelationService and add TableStorageService * Refactor RelationService to implement IRelationService
1 parent 0aea321 commit b89bd66

File tree

5 files changed

+161
-38
lines changed

5 files changed

+161
-38
lines changed

‎CentralMemoryMcp.Functions/Functions/GraphFunctions.cs‎

Lines changed: 9 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -2,20 +2,13 @@
22
using Microsoft.Azure.Functions.Worker;
33
using Microsoft.Azure.Functions.Worker.Extensions.Mcp;
44
using System.ComponentModel;
5-
using System; // for Guid
5+
using CentralMemoryMcp.Functions.Services;
66

77
namespace CentralMemoryMcp.Functions.Functions;
88

9-
public class GraphFunctions
9+
public class GraphFunctions(IKnowledgeGraphService graph, IRelationService relations)
1010
{
11-
private readonly IKnowledgeGraphService _graph;
12-
private readonly IRelationService _relations;
13-
14-
public GraphFunctions(IKnowledgeGraphService graph, IRelationService relations)
15-
{
16-
_graph = graph;
17-
_relations = relations;
18-
}
11+
private readonly IRelationService _relations = relations;
1912

2013
[Function(nameof(ReadGraph))]
2114
public async Task<object> ReadGraph(
@@ -24,7 +17,7 @@ public async Task<object> ReadGraph(
2417
[McpToolProperty("workspaceName", "The unique identifier of the workspace.", isRequired: true)]
2518
string workspaceName)
2619
{
27-
var entities = await _graph.ReadGraphAsync(workspaceName);
20+
var entities = await graph.ReadGraphAsync(workspaceName);
2821
var relations = await _relations.GetRelationsForWorkspaceAsync(workspaceName);
2922

3023
return new
@@ -55,7 +48,7 @@ public async Task<object> UpsertEntity(
5548
request.Observations ?? [],
5649
request.Metadata);
5750

58-
model = await _graph.UpsertEntityAsync(model); // capture potentially reused Id
51+
model = await graph.UpsertEntityAsync(model); // capture potentially reused Id
5952
return new { success = true, id = model.Id, workspace = model.WorkspaceName, name = model.Name };
6053
}
6154

@@ -80,12 +73,12 @@ public async Task<object> UpsertRelation(
8073

8174
if (fromId == Guid.Empty && !string.IsNullOrWhiteSpace(request.From))
8275
{
83-
var entity = await _graph.GetEntityAsync(request.WorkspaceName, request.From);
76+
var entity = await graph.GetEntityAsync(request.WorkspaceName, request.From);
8477
if (entity is not null) fromId = entity.Id; else return new { success = false, message = $"Source entity '{request.From}' not found." };
8578
}
8679
if (toId == Guid.Empty && !string.IsNullOrWhiteSpace(request.To))
8780
{
88-
var entity = await _graph.GetEntityAsync(request.WorkspaceName, request.To);
81+
var entity = await graph.GetEntityAsync(request.WorkspaceName, request.To);
8982
if (entity is not null) toId = entity.Id; else return new { success = false, message = $"Target entity '{request.To}' not found." };
9083
}
9184

@@ -116,14 +109,14 @@ public async Task<object> GetEntityRelations(
116109
[McpToolProperty("entityName", "Legacy entity name (used if entityId not provided).", isRequired: false)]
117110
string? entityName)
118111
{
119-
Guid resolvedId = Guid.Empty;
112+
Guid resolvedId;
120113
if (entityId.HasValue && entityId.Value != Guid.Empty)
121114
{
122115
resolvedId = entityId.Value;
123116
}
124117
else if (!string.IsNullOrWhiteSpace(entityName))
125118
{
126-
var entity = await _graph.GetEntityAsync(workspaceName, entityName);
119+
var entity = await graph.GetEntityAsync(workspaceName, entityName);
127120
if (entity is null)
128121
{
129122
return new { success = false, message = $"Entity '{entityName}' not found in workspace '{workspaceName}'." };

‎CentralMemoryMcp.Functions/Program.cs‎

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using Azure.Data.Tables;
2-
using CentralMemoryMcp.Functions;
2+
using CentralMemoryMcp.Functions.Services;
3+
using CentralMemoryMcp.Functions.Storage;
34
using Microsoft.Azure.Functions.Worker.Builder;
45
using Microsoft.Extensions.DependencyInjection;
56
using Microsoft.Extensions.Hosting;
@@ -10,14 +11,6 @@
1011
// Services
1112
builder.ConfigureFunctionsWebApplication();
1213

13-
// Register Azure Table Storage
14-
builder.Services.AddSingleton(sp =>
15-
{
16-
var connectionString = Environment.GetEnvironmentVariable("AzureWebJobsStorage")
17-
?? throw new InvalidOperationException("AzureWebJobsStorage connection string is not configured.");
18-
return new TableServiceClient(connectionString);
19-
});
20-
2114
// Register application services
2215
builder.Services.AddSingleton<ITableStorageService, TableStorageService>();
2316
builder.Services.AddSingleton<IKnowledgeGraphService, KnowledgeGraphService>();

‎CentralMemoryMcp.Functions/Services/KnowledgeGraphService.cs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
using Azure;
22
using Azure.Data.Tables;
33
using CentralMemoryMcp.Functions.Models;
4+
using CentralMemoryMcp.Functions.Storage;
45

5-
namespace CentralMemoryMcp.Functions;
6+
namespace CentralMemoryMcp.Functions.Services;
67

78
public interface IKnowledgeGraphService
89
{

‎CentralMemoryMcp.Functions/Services/RelationService.cs‎

Lines changed: 133 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
using Azure;
22
using Azure.Data.Tables;
33
using CentralMemoryMcp.Functions.Models;
4+
using CentralMemoryMcp.Functions.Storage;
45

5-
namespace CentralMemoryMcp.Functions;
6+
namespace CentralMemoryMcp.Functions.Services;
67

78
public interface IRelationService
89
{
@@ -19,7 +20,7 @@ public async Task<RelationModel> UpsertRelationAsync(RelationModel model, Cancel
1920
{
2021
var table = await storage.GetRelationsTableAsync(ct);
2122
// Check for existing relation (same workspace, from, to, type)
22-
string filter = $"PartitionKey eq '{model.WorkspaceName}' and FromEntityId eq '{model.FromEntityId.ToString("N")}' and ToEntityId eq '{model.ToEntityId.ToString("N")}' and RelationType eq '{EscapeFilterValue(model.RelationType)}'";
23+
var filter = $"PartitionKey eq '{model.WorkspaceName}' and FromEntityId eq '{model.FromEntityId:N}' and ToEntityId eq '{model.ToEntityId:N}' and RelationType eq '{EscapeFilterValue(model.RelationType)}'";
2324
await foreach(var e in table.QueryAsync<TableEntity>(filter: filter, maxPerPage:1, cancellationToken: ct))
2425
{
2526
// Reuse its Id
@@ -47,17 +48,142 @@ public async Task<RelationModel> UpsertRelationAsync(RelationModel model, Cancel
4748
var table = await storage.GetRelationsTableAsync(ct);
4849
try
4950
{
50-
var partitionKey = workspaceName;
51-
var rowKey = relationId.ToString("N");
52-
var response = await table.GetEntityAsync<TableEntity>(partitionKey, rowKey, cancellationToken: ct);
53-
var e = response.Value;
51+
var response = await table.GetEntityAsync<TableEntity>(workspaceName, relationId.ToString("N"), cancellationToken: ct);
52+
var model = new RelationModel(
53+
response.Value.GetString("WorkspaceName")!,
54+
Guid.Parse(response.Value.GetString("FromEntityId")!),
55+
Guid.Parse(response.Value.GetString("ToEntityId")!),
56+
response.Value.GetString("RelationType")!,
57+
response.Value.GetString("Metadata"))
58+
{
59+
Id = relationId
60+
};
61+
return model;
62+
}
63+
catch (RequestFailedException ex) when (ex.Status == 404)
64+
{
65+
return null;
66+
}
67+
}
68+
69+
public async Task<List<RelationModel>> GetRelationsFromEntityAsync(string workspaceName, Guid fromEntityId, CancellationToken ct = default)
70+
{
71+
var table = await storage.GetRelationsTableAsync(ct);
72+
var results = new List<RelationModel>();
73+
var fromIdStr = fromEntityId.ToString("N");
74+
await foreach (var e in table.QueryAsync<TableEntity>(
75+
filter: $"PartitionKey eq '{workspaceName}' and FromEntityId eq '{fromIdStr}'",
76+
cancellationToken: ct))
77+
{
78+
var relationId = Guid.TryParse(e.GetString("Id"), out var rid) ? rid : Guid.NewGuid();
5479
var model = new RelationModel(
5580
e.GetString("WorkspaceName")!,
5681
Guid.Parse(e.GetString("FromEntityId")!),
5782
Guid.Parse(e.GetString("ToEntityId")!),
5883
e.GetString("RelationType")!,
5984
e.GetString("Metadata"));
6085
model.Id = relationId;
86+
results.Add(model);
87+
}
88+
return results;
89+
}
90+
91+
public async Task<List<RelationModel>> GetRelationsForWorkspaceAsync(string workspaceName, CancellationToken ct = default)
92+
{
93+
var table = await storage.GetRelationsTableAsync(ct);
94+
var results = new List<RelationModel>();
95+
await foreach (var e in table.QueryAsync<TableEntity>(
96+
filter: $"PartitionKey eq '{workspaceName}'",
97+
cancellationToken: ct))
98+
{
99+
var relationId = Guid.TryParse(e.GetString("Id"), out var rid) ? rid : Guid.NewGuid();
100+
var model = new RelationModel(
101+
e.GetString("WorkspaceName")!,
102+
Guid.Parse(e.GetString("FromEntityId")!),
103+
Guid.Parse(e.GetString("ToEntityId")!),
104+
e.GetString("RelationType")!,
105+
e.GetString("Metadata"));
106+
model.Id = relationId;
107+
results.Add(model);
108+
}
109+
return results;
110+
}
111+
112+
public async Task DeleteRelationAsync(string workspaceName, Guid relationId, CancellationToken ct = default)
113+
{
114+
var table = await storage.GetRelationsTableAsync(ct);
115+
try
116+
{
117+
await table.DeleteEntityAsync(workspaceName, relationId.ToString("N"), cancellationToken: ct);
118+
}
119+
catch (RequestFailedException ex) when (ex.Status == 404)
120+
{
121+
// not found; ignore
122+
}
123+
}
124+
125+
private static string EscapeFilterValue(string value) => value.Replace("'", "''");
126+
}
127+
using Azure;
128+
using Azure.Data.Tables;
129+
using CentralMemoryMcp.Functions.Models;
130+
using CentralMemoryMcp.Functions.Storage;
131+
132+
namespace CentralMemoryMcp.Functions.Services;
133+
134+
public interface IRelationService
135+
{
136+
Task<RelationModel> UpsertRelationAsync(RelationModel model, CancellationToken ct = default);
137+
Task<RelationModel?> GetRelationAsync(string workspaceName, Guid relationId, CancellationToken ct = default);
138+
Task<List<RelationModel>> GetRelationsFromEntityAsync(string workspaceName, Guid fromEntityId, CancellationToken ct = default);
139+
Task<List<RelationModel>> GetRelationsForWorkspaceAsync(string workspaceName, CancellationToken ct = default);
140+
Task DeleteRelationAsync(string workspaceName, Guid relationId, CancellationToken ct = default);
141+
}
142+
143+
public class RelationService(ITableStorageService storage) : IRelationService
144+
{
145+
public async Task<RelationModel> UpsertRelationAsync(RelationModel model, CancellationToken ct = default)
146+
{
147+
var table = await storage.GetRelationsTableAsync(ct);
148+
// Check for existing relation (same workspace, from, to, type)
149+
var filter = $"PartitionKey eq '{model.WorkspaceName}' and FromEntityId eq '{model.FromEntityId:N}' and ToEntityId eq '{model.ToEntityId:N}' and RelationType eq '{EscapeFilterValue(model.RelationType)}'";
150+
await foreach(var e in table.QueryAsync<TableEntity>(filter: filter, maxPerPage:1, cancellationToken: ct))
151+
{
152+
// Reuse its Id
153+
if (e.TryGetValue("Id", out var idObj) && idObj is string idStr && Guid.TryParse(idStr, out var existingId))
154+
{
155+
model.Id = existingId;
156+
}
157+
break;
158+
}
159+
var entity = new TableEntity(model.PartitionKey, model.RowKey)
160+
{
161+
{"Id", model.Id.ToString("N")},
162+
{"WorkspaceName", model.WorkspaceName},
163+
{"FromEntityId", model.FromEntityId.ToString("N")},
164+
{"ToEntityId", model.ToEntityId.ToString("N")},
165+
{"RelationType", model.RelationType},
166+
{"Metadata", model.Metadata ?? string.Empty}
167+
};
168+
await table.UpsertEntityAsync(entity, TableUpdateMode.Replace, ct);
169+
return model;
170+
}
171+
172+
public async Task<RelationModel?> GetRelationAsync(string workspaceName, Guid relationId, CancellationToken ct = default)
173+
{
174+
var table = await storage.GetRelationsTableAsync(ct);
175+
try
176+
{
177+
var response = await table.GetEntityAsync<TableEntity>(workspaceName, relationId.ToString("N"), cancellationToken: ct);
178+
var model = new RelationModel(
179+
response.Value.GetString("WorkspaceName")!,
180+
Guid.Parse(response.Value.GetString("FromEntityId")!),
181+
Guid.Parse(response.Value.GetString("ToEntityId")!),
182+
response.Value.GetString("RelationType")!,
183+
response.Value.GetString("Metadata"))
184+
{
185+
Id = relationId
186+
};
61187
return model;
62188
}
63189
catch (RequestFailedException ex) when (ex.Status == 404)
@@ -112,11 +238,9 @@ public async Task<List<RelationModel>> GetRelationsForWorkspaceAsync(string work
112238
public async Task DeleteRelationAsync(string workspaceName, Guid relationId, CancellationToken ct = default)
113239
{
114240
var table = await storage.GetRelationsTableAsync(ct);
115-
var partitionKey = workspaceName;
116-
var rowKey = relationId.ToString("N");
117241
try
118242
{
119-
await table.DeleteEntityAsync(partitionKey, rowKey, cancellationToken: ct);
243+
await table.DeleteEntityAsync(workspaceName, relationId.ToString("N"), cancellationToken: ct);
120244
}
121245
catch (RequestFailedException ex) when (ex.Status == 404)
122246
{

‎CentralMemoryMcp.Functions/Storage/TableStorageService.cs‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
using Azure;
1+
using Azure.Identity;
22
using Azure.Data.Tables;
33

4-
namespace CentralMemoryMcp.Functions
4+
namespace CentralMemoryMcp.Functions.Storage
55
{
66
public interface ITableStorageService
77
{
@@ -17,7 +17,19 @@ public class TableStorageService : ITableStorageService
1717
private const string RelationsTableName = "relations";
1818
private const string WorkspacesTableName = "workspaces";
1919

20-
public TableStorageService(TableServiceClient serviceClient) => _serviceClient = serviceClient;
20+
public TableStorageService()
21+
{
22+
var conn = Environment.GetEnvironmentVariable("AzureWebJobsStorage");
23+
if (!string.IsNullOrWhiteSpace(conn))
24+
{
25+
_serviceClient = new TableServiceClient(conn);
26+
return;
27+
}
28+
29+
var endpoint = Environment.GetEnvironmentVariable("AzureWebJobsStorage__tableServiceUri")
30+
?? throw new InvalidOperationException("AzureWebJobsStorage is not configured for managed identity.");
31+
_serviceClient = new TableServiceClient(new Uri(endpoint), new DefaultAzureCredential());
32+
}
2133

2234
public async Task<TableClient> GetEntitiesTableAsync(CancellationToken ct = default)
2335
{

0 commit comments

Comments
 (0)