Skip to content

Commit

Permalink
RavenDB-22741 : Add support for querying by prefix
Browse files Browse the repository at this point in the history
  • Loading branch information
aviv committed Aug 11, 2024
1 parent 8270d29 commit 94622c6
Show file tree
Hide file tree
Showing 6 changed files with 460 additions and 19 deletions.
6 changes: 6 additions & 0 deletions src/Raven.Client/Constants.cs
Original file line number Diff line number Diff line change
Expand Up @@ -442,6 +442,12 @@ private Sharding()
}

internal const string ShardContextParameterName = "__shardContext";

internal const string ShardContextDocumentIds = "DocumentIds";

internal const string ShardContextPrefixes = "Prefixes";


}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,10 @@ protected void ShardContext(Action<IQueryShardedContextBuilder> builder)

builder.Invoke(builderImpl);

object shardContext;

if (builderImpl.DocumentIds.Count == 1)
shardContext = builderImpl.DocumentIds.First();
else
shardContext = builderImpl.DocumentIds;

QueryParameters.Add(Constants.Documents.Querying.Sharding.ShardContextParameterName, shardContext);
QueryParameters.Add(Constants.Documents.Querying.Sharding.ShardContextParameterName, new
{
builderImpl.DocumentIds,
builderImpl.Prefixes
});
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -7,4 +7,10 @@ public interface IQueryShardedContextBuilder
IQueryShardedContextBuilder ByDocumentId(string id);

IQueryShardedContextBuilder ByDocumentIds(IEnumerable<string> ids);

IQueryShardedContextBuilder ByPrefix(string prefix);

IQueryShardedContextBuilder ByPrefixes(IEnumerable<string> prefixes);


}
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@ internal sealed class QueryShardedContextBuilder : IQueryShardedContextBuilder
{
public HashSet<string> DocumentIds { get; } = new(StringComparer.OrdinalIgnoreCase);

public HashSet<string> Prefixes { get; } = new(StringComparer.OrdinalIgnoreCase);


public IQueryShardedContextBuilder ByDocumentId(string id)
{
DocumentIds.Add(id);
Expand All @@ -23,4 +26,20 @@ public IQueryShardedContextBuilder ByDocumentIds(IEnumerable<string> ids)

return this;
}

public IQueryShardedContextBuilder ByPrefix(string prefix)
{
Prefixes.Add(prefix);
return this;
}

public IQueryShardedContextBuilder ByPrefixes(IEnumerable<string> prefixes)
{
foreach (string prefix in prefixes)
{
Prefixes.Add(prefix);
}

return this;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
using Raven.Client.Exceptions.Sharding;
using Raven.Client.Extensions;
using Raven.Client.Http;
using Raven.Client.ServerWide.Sharding;
using Raven.Client.Util;
using Raven.Server.Documents.Indexes;
using Raven.Server.Documents.Indexes.MapReduce.Auto;
Expand All @@ -25,6 +26,7 @@
using Raven.Server.Documents.Sharding.Handlers;
using Raven.Server.Documents.Sharding.Operations;
using Raven.Server.ServerWide.Context;
using Raven.Server.ServerWide.Sharding;
using Sparrow;
using Sparrow.Json;
using Sparrow.Json.Parsing;
Expand All @@ -42,7 +44,7 @@ public abstract class AbstractShardedQueryProcessor<TCommand, TResult, TCombined
private readonly bool _indexEntriesOnly;
private readonly bool _ignoreLimit;
private readonly string _raftUniqueRequestId;
private readonly HashSet<int> _filteredShardIndexes;
private HashSet<int> _filteredShardIndexes;

protected readonly TransactionOperationContext Context;
protected readonly ShardedDatabaseRequestHandler RequestHandler;
Expand Down Expand Up @@ -87,35 +89,74 @@ protected AbstractShardedQueryProcessor(
_raftUniqueRequestId = RequestHandler.GetRaftRequestIdFromQuery() ?? RaftIdGenerator.NewId();

if (Query.QueryParameters != null && Query.QueryParameters.TryGetMember(Constants.Documents.Querying.Sharding.ShardContextParameterName, out object filter))
HandleShardsFiltering(context, filter);
}

private void HandleShardsFiltering(TransactionOperationContext context, object filter)
{
_filteredShardIndexes = new HashSet<int>();

if (filter is not BlittableJsonReaderObject filterBlittable)
{
// v6.0
// User can define a query parameter ("__shardContext") which is an id or an array
// that contains the ids whose shards the query should be limited to.
// Advanced: Optimization if user wants to run a query and knows what shards it is on. Such as:
// from Orders where State = $state and User = $user where all the orders are on the same share as the user

_filteredShardIndexes = new HashSet<int>();
switch (filter)
{
case LazyStringValue:
case LazyCompressedStringValue:
_filteredShardIndexes.Add(RequestHandler.DatabaseContext.GetShardNumberFor(context, filter.ToString()));
break;
case BlittableJsonReaderArray arr:
{
for (int i = 0; i < arr.Length; i++)
{
var it = arr.GetStringByIndex(i);
_filteredShardIndexes.Add(RequestHandler.DatabaseContext.GetShardNumberFor(context, it));
}
break;
}
GetShardNumbersForFiltering(context, arr);
break;
default:
throw new NotSupportedException($"Unknown type of a shard context query parameter: {filter.GetType().Name}");
}
}

else
{
_filteredShardIndexes = null;
// v6.1
// "__shardContext" can also be an object, with DocumentIds and Prefixes fields
// now, in addition to document ids, users can also specify prefixes to query by

var prefixedSetting = RequestHandler.DatabaseContext.DatabaseRecord.Sharding.Prefixed;

if (prefixedSetting.Count > 0 &&
filterBlittable.TryGetMember(Constants.Documents.Querying.Sharding.ShardContextPrefixes, out var p) &&
p is BlittableJsonReaderArray { Length: > 0 } prefixesArr)
{
foreach (var item in prefixesArr)
{
var location = prefixedSetting.BinarySearch(new PrefixedShardingSetting(item.ToString()), PrefixedSettingComparer.Instance);
if (location < 0)
throw new InvalidOperationException($"Prefix '{item}' wasn't found in Sharding Configuration");

foreach (var shard in prefixedSetting[location].Shards)
{
_filteredShardIndexes.Add(shard);
}
}
}

if (filterBlittable.TryGetMember(Constants.Documents.Querying.Sharding.ShardContextDocumentIds, out var o) &&
o is BlittableJsonReaderArray { Length: > 0 } idsArr)
{
GetShardNumbersForFiltering(context, idsArr);
}
}
}

private void GetShardNumbersForFiltering(TransactionOperationContext context, BlittableJsonReaderArray array)
{
for (int i = 0; i < array.Length; i++)
{
var it = array.GetStringByIndex(i);
_filteredShardIndexes.Add(RequestHandler.DatabaseContext.GetShardNumberFor(context, it));
}
}

Expand Down
Loading

0 comments on commit 94622c6

Please sign in to comment.