Query: Fixes default to BadRequestException in case of internal errors in ServiceInterop (#3399)

* Don't default to BadRequestException in case of errors in ServiceInterop

* Incorporate code review feedback

* Fix build error

* fix up failing test

Author: neildsh

Microsoft.Azure.Cosmos/src/Query/Core/ExceptionToCosmosException.cs

@@ -76,32 +76,28 @@ private static bool TryCreateFromExceptionWithStackTrace(
             ITrace trace,
             out CosmosException cosmosException)
         {
-            // Use the original stack trace from the inner exception.
-            if (exceptionWithStackTrace.InnerException is Microsoft.Azure.Documents.DocumentClientException
-                || exceptionWithStackTrace.InnerException is CosmosException)
-            {
-                return ExceptionToCosmosException.TryCreateFromException(
-                    exceptionWithStackTrace.InnerException, 
-                    trace, 
-                    out cosmosException);
-            }
+            Exception innerException = ExceptionWithStackTraceException.UnWrapMonadExcepion(exceptionWithStackTrace, trace);
 
             if (!ExceptionToCosmosException.TryCreateFromException(
-                exceptionWithStackTrace.InnerException,
+                innerException,
                 trace,
                 out cosmosException))
             {
                 return false;
             }
 
-            cosmosException = CosmosExceptionFactory.Create(
-                cosmosException.StatusCode,
-                cosmosException.Message,
-                exceptionWithStackTrace.StackTrace,
-                headers: cosmosException.Headers,
-                cosmosException.Trace,
-                cosmosException.Error,
-                cosmosException.InnerException);
+            if (innerException is not CosmosException && innerException is not Documents.DocumentClientException)
+            {
+                cosmosException = CosmosExceptionFactory.Create(
+                    cosmosException.StatusCode,
+                    cosmosException.Message,
+                    exceptionWithStackTrace.StackTrace,
+                    headers: cosmosException.Headers,
+                    cosmosException.Trace,
+                    cosmosException.Error,
+                    cosmosException.InnerException);
+            }
+
             return true;
         }
 

Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs

@@ -76,18 +76,14 @@ public static async Task<PartitionedQueryExecutionInfo> GetQueryPlanWithServiceI
 
                 if (!tryGetQueryPlan.Succeeded)
                 {
-                    Exception originalException = ExceptionWithStackTraceException.UnWrapMonadExcepion(tryGetQueryPlan.Exception, serviceInteropTrace);
-                    if (originalException is CosmosException)
+                    if (ExceptionToCosmosException.TryCreateFromException(tryGetQueryPlan.Exception, serviceInteropTrace, out CosmosException cosmosException))
                     {
-                        throw originalException;
+                        throw cosmosException;
+                    }
+                    else
+                    {
+                        throw ExceptionWithStackTraceException.UnWrapMonadExcepion(tryGetQueryPlan.Exception, serviceInteropTrace);
                     }
-
-                    throw CosmosExceptionFactory.CreateBadRequestException(
-                        message: originalException.Message,
-                        headers: new Headers(),
-                        stackTrace: tryGetQueryPlan.Exception.StackTrace,
-                        innerException: originalException,
-                        trace: trace);
                 }
 
                 return tryGetQueryPlan.Result;

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanRetrieverTests.cs

@@ -41,7 +41,6 @@ public async Task ServiceInterop_BadRequestContainsInnerException()
                 It.IsAny<Cosmos.GeospatialType>(),
                 It.IsAny<CancellationToken>())).ReturnsAsync(TryCatch<PartitionedQueryExecutionInfo>.FromException(innerException));
 
-            Mock<ITrace> trace = new Mock<ITrace>();
             CosmosException cosmosException = await Assert.ThrowsExceptionAsync<CosmosException>(() => QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync(
                 queryClient.Object,
                 new SqlQuerySpec("selectttttt * from c"),
@@ -50,8 +49,7 @@ public async Task ServiceInterop_BadRequestContainsInnerException()
                 hasLogicalPartitionKey: false,
                 geospatialType: Cosmos.GeospatialType.Geography,
                 useSystemPrefix: false,
-                trace.Object,
-                default));
+                NoOpTrace.Singleton));
 
             Assert.AreEqual(HttpStatusCode.BadRequest, cosmosException.StatusCode);
             Assert.AreEqual(innerException, cosmosException.InnerException);
@@ -92,5 +90,40 @@ public async Task ServiceInterop_BadRequestContainsOriginalCosmosException()
 
             Assert.AreEqual(expectedException, cosmosException);
         }
+
+        [TestMethod]
+        public async Task ServiceInterop_E_UNEXPECTED()
+        {
+            UnexpectedQueryPartitionProviderException innerException = new UnexpectedQueryPartitionProviderException("E_UNEXPECTED");
+            Mock<CosmosQueryClient> queryClient = new Mock<CosmosQueryClient>();
+
+            queryClient.Setup(c => c.TryGetPartitionedQueryExecutionInfoAsync(
+                It.IsAny<SqlQuerySpec>(),
+                It.IsAny<ResourceType>(),
+                It.IsAny<Documents.PartitionKeyDefinition>(),
+                It.IsAny<bool>(),
+                It.IsAny<bool>(),
+                It.IsAny<bool>(),
+                It.IsAny<bool>(),
+                It.IsAny<bool>(),
+                It.IsAny<bool>(),
+                It.IsAny<Cosmos.GeospatialType>(),
+                It.IsAny<CancellationToken>())).ReturnsAsync(TryCatch<PartitionedQueryExecutionInfo>.FromException(innerException));
+
+            CosmosException cosmosException = await Assert.ThrowsExceptionAsync<CosmosException>(() => QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync(
+                queryClient.Object,
+                new SqlQuerySpec("Super secret query that triggers bug"),
+                ResourceType.Document,
+                new Documents.PartitionKeyDefinition() { Paths = new Collection<string>() { "/id" } },
+                hasLogicalPartitionKey: false,
+                geospatialType: Cosmos.GeospatialType.Geography,
+                useSystemPrefix: false,
+                NoOpTrace.Singleton));
+
+            Assert.AreEqual(HttpStatusCode.InternalServerError, cosmosException.StatusCode);
+            Assert.AreEqual(innerException, cosmosException.InnerException);
+            Assert.IsNotNull(cosmosException.Trace);
+            Assert.IsNotNull(cosmosException.Diagnostics);
+        }
     }
 }