[Internal] Query: Adds Split Support for Ode

Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs

@@ -126,6 +126,38 @@ public async Task<TryCatch<T>> TryAsync<T>(
  return matchResult;
  }
 
+ public async ValueTask<TryCatch<T>> TryAsync<T>(
+ Func<TResult, ValueTask<T>> onSuccess)
+ {
+ TryCatch<T> matchResult;
+ if (this.Succeeded)
+ {
+ matchResult = TryCatch<T>.FromResult(await onSuccess(this.either.FromRight(default)));
+ }
+ else
+ {
+ matchResult = TryCatch<T>.FromException(this.either.FromLeft(default));
+ }
+
+ return matchResult;
+ }
+
+ public TryCatch<T> Try<T>(
+ Func<TResult, TryCatch<T>> onSuccess)
+ {
+ TryCatch<T> matchResult;
+ if (this.Succeeded)
+ {
+ matchResult = onSuccess(this.either.FromRight(default));
+ }
+ else
+ {
+ matchResult = TryCatch<T>.FromException(this.either.FromLeft(default));
+ }
+
+ return matchResult;
+ }
+
  public TryCatch<TResult> Catch(
  Action<Exception> onError)
  {

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs

@@ -0,0 +1,28 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// ------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline
+{
+ using System;
+
+ internal static class CosmosExceptionExtensions
+ {
+ public static bool IsPartitionSplitException(this Exception ex)
+ {
+ if (ex != null)
+ {
+ return IsPartitionSplitException(ex as CosmosException);
+ }
+
+ return false;
+ }
+
+ public static bool IsPartitionSplitException(this CosmosException ex)
+ {
+ return ex is CosmosException cosmosException
+ && (cosmosException.StatusCode == System.Net.HttpStatusCode.Gone)
+ && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone);
+ }
+ }
+}

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/OrderByCrossPartitionQueryPipelineStage.cs

@@ -1153,9 +1153,7 @@ private static bool IsSplitException(Exception exception)
  exception = exception.InnerException;
  }
 
- return exception is CosmosException cosmosException
- && (cosmosException.StatusCode == HttpStatusCode.Gone)
- && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone);
+ return exception.IsPartitionSplitException();
  }
 
  public void SetCancellationToken(CancellationToken cancellationToken)

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionContinuationToken.cs

@@ -43,11 +43,11 @@ public static CosmosElement ToCosmosElement(OptimisticDirectExecutionContinuatio
  public static TryCatch<OptimisticDirectExecutionContinuationToken> TryCreateFromCosmosElement(CosmosElement cosmosElement)
  {
  CosmosObject cosmosObjectContinuationToken = cosmosElement as CosmosObject;
- if (cosmosObjectContinuationToken == null)
+ if (cosmosObjectContinuationToken == null || !cosmosObjectContinuationToken.ContainsKey(OptimisticDirectExecutionToken))
  {
  return TryCatch<OptimisticDirectExecutionContinuationToken>.FromException(
- new MalformedChangeFeedContinuationTokenException(
- message: $"Malformed Continuation Token"));
+ new MalformedContinuationTokenException(
+ message: $"Malformed Continuation Token: Expected OptimisticDirectExecutionToken\r\n"));
  }
 
  TryCatch<ParallelContinuationToken> inner = ParallelContinuationToken.TryCreateFromCosmosElement(cosmosObjectContinuationToken[OptimisticDirectExecutionToken]);