Support hierarchical partition keys in Cosmos DB grain storage
- Dominant language
- C#
- Stars
- 10.9k
- Forks
- 2.1k
- Avg merge
- 14h 42m
- Merged PRs (30d)
- 354
Description
## Summary
The current `IPartitionKeyProvider` interface returns a `string`, which limits Cosmos DB grain storage to single-level partition keys. With [hierarchical partition keys](https://learn.microsoft.com/en-us/azure/cosmos-db/hierarchical-partition-keys) now GA in Cosmos DB, it would be valuable to support multi-level partitioning for Orleans grain state.
## Background
I'm returning to Orleans after a hiatus and starting a new project that would benefit from hierarchical partitioning. The current implementation wraps the string result in `new PartitionKey(partitionKey)`:
```csharp
// Current implementation in CosmosGrainStorage.cs
private ValueTask BuildPartitionKey(string grainType, GrainId grainId) =>
_partitionKeyProvider.GetPartitionKey(grainType, grainId);
// Later used as:
var pk = new PartitionKey(partitionKey);
```
## Proposal
Add a virtual default interface method (DIM) to `IPartitionKeyProvider` that returns a `PartitionKey` directly, allowing implementations to construct hierarchical keys:
```csharp
public interface IPartitionKeyProvider
{
// Existing method (unchanged for backwards compatibility)
ValueTask GetPartitionKey(string grainType, GrainId grainId);
// New DIM for hierarchical partition key support
ValueTask GetPartitionKeyValue(string grainType, GrainId grainId)
=> new(new PartitionKey(GetPartitionKey(grainType, grainId).Result));
}
```
Then update `CosmosGrainStorage` to call `GetPartitionKeyValue()` instead of wrapping the string manually.
### Example Usage
```csharp
public class TenantAwarePartitionKeyProvider : IPartitionKeyProvider
{
public ValueTask GetPartitionKey(string grainType, GrainId grainId)
=> new(grainType); // Legacy fallback
public ValueTask GetPartitionKeyValue(string grainType, GrainId grainId)
{
var tenantId = ExtractTenantFromGrainId(grainId);
// Hierarchical: TenantId -> GrainType -> Region
return new(new PartitionKeyBuilder()
.Add(tenantId)
.Add(grainType)
.Build());
}
}
```
## Benefits
1. **Multi-tenant scenarios**: Partition first by tenant, then by grain type
2. **Better data distribution**: Avoid hot partitions in high-volume scenarios
3. **Backwards compatible**: Existing implementations continue to work unchanged via the DIM
4. **Cosmos DB best practices**: Aligns with [Microsoft's guidance on partition key design](https://learn.microsoft.com/en-us/azure/cosmos-db/partitioning-overview#choose-a-partition-key)
## Considerations
- `CosmosStorageOptions.PartitionKeyPath` would need to support hierarchical paths (e.g., `/TenantId/GrainType`) (we should be able to define GrainStateEntity as an `object` w/ nested properties rather than a string)
- Container creation logic may need updates for hierarchical key configuration
- Documentation updates for multi-level partitioning patterns
## Environment
- Orleans version: 10.x
- Target: `Orleans.Persistence.Cosmos`
Happy to submit a PR if this direction makes sense.
Contributor guide
Research direction
Start with IPartitionKeyProvider and CosmosGrainStorage.cs to trace the current string partition-key flow and the proposed direct PartitionKey path. Then inspect CosmosStorageOptions.PartitionKeyPath, GrainStateEntity, and container creation logic for hierarchical configuration. Done means supporting hierarchical keys while preserving existing providers, with the related configuration and documentation updates identified in the issue.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- azure, csharp
- Domain
- databases, distributed-systems
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100