Added IsPartitionSplitException() overload to take CosmosElement

Azure · Jan 13, 2023 · 2416a89 · 2416a89
1 parent c94db62
commit 2416a89
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 18 deletions.
diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosExceptionExtensions.cs
@@ -9,10 +9,20 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline
  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);
+ && (cosmosException.StatusCode == System.Net.HttpStatusCode.Gone)
+ && (cosmosException.SubStatusCode == (int)Documents.SubStatusCodes.PartitionKeyRangeGone);
  }
  }
 }
diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs
@@ -305,19 +305,22 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateFromPartitione
  trace,
  cancellationToken);
  }
- else if (createPassthroughQuery)
- {
- SetTestInjectionPipelineType(inputParameters, Passthrough);
-
- tryCreatePipelineStage = CosmosQueryExecutionContextFactory.TryCreatePassthroughQueryExecutionContext(
- documentContainer,
- inputParameters,
- targetRanges,
- cancellationToken);
- }
  else
  {
- tryCreatePipelineStage = TryCreateSpecializedDocumentQueryExecutionContext(documentContainer, cosmosQueryContext, inputParameters, targetRanges, partitionedQueryExecutionInfo, cancellationToken);
+ if (createPassthroughQuery)
+ {
+ SetTestInjectionPipelineType(inputParameters, Passthrough);
+
+ tryCreatePipelineStage = CosmosQueryExecutionContextFactory.TryCreatePassthroughQueryExecutionContext(
+ documentContainer,
+ inputParameters,
+ targetRanges,
+ cancellationToken);
+ }
+ else
+ {
+ tryCreatePipelineStage = TryCreateSpecializedDocumentQueryExecutionContext(documentContainer, cosmosQueryContext, inputParameters, targetRanges, partitionedQueryExecutionInfo, cancellationToken);
+ }
  }
 
  return tryCreatePipelineStage;

diff --git a/...s/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/...s/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs
@@ -623,7 +623,8 @@ public async Task<Exception> ShouldReturnFailure()
  activityId: "0f8fad5b-d9cb-469f-a165-70867728950e",
  requestCharge: default);
  }
- else if (this.IsFailedFallbackPipelineTest && this.GoneExceptionCreated && !this.TooManyRequestsFailureCreated)
+
+ if (this.IsFailedFallbackPipelineTest && this.GoneExceptionCreated && !this.TooManyRequestsFailureCreated)
  {
  this.TooManyRequestsFailureCreated = true;
  return new CosmosException(
@@ -633,10 +634,8 @@ public async Task<Exception> ShouldReturnFailure()
  activityId: "111fad5b-d9cb-469f-a165-70867728950e",
  requestCharge: 0);
  }
- else
- {
- return null;
- }
+
+ return null;
  }
  }
  }