internal TryCatch TryGetPartitionedQueryExecutionInfoInternal()

in Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs [185:336]
134 lines of code
12 McCabe index (conditional complexity)

        internal TryCatch<PartitionedQueryExecutionInfoInternal> TryGetPartitionedQueryExecutionInfoInternal(
            string querySpecJsonString,
            PartitionKeyDefinition partitionKeyDefinition,
            VectorEmbeddingPolicy vectorEmbeddingPolicy,
            bool requireFormattableOrderByQuery,
            bool isContinuationExpected,
            bool allowNonValueAggregateQuery,
            bool hasLogicalPartitionKey,
            bool allowDCount,
            bool useSystemPrefix,
            bool hybridSearchSkipOrderByRewrite,
            GeospatialType geospatialType)
        {
            if (querySpecJsonString == null || partitionKeyDefinition == null)
            {
                return TryCatch<PartitionedQueryExecutionInfoInternal>.FromResult(DefaultInfoInternal);
            }

            List<string> paths = new List<string>(partitionKeyDefinition.Paths);
            List<IReadOnlyList<string>> pathPartsList = new List<IReadOnlyList<string>>(paths.Count);
            uint[] partsLengths = new uint[paths.Count];
            int allPartsLength = 0;

            for (int i = 0; i < paths.Count; i++)
            {
                IReadOnlyList<string> pathParts = PathParser.GetPathParts(paths[i]);
                partsLengths[i] = (uint)pathParts.Count;
                pathPartsList.Add(pathParts);
                allPartsLength += pathParts.Count;
            }

            string[] allParts = new string[allPartsLength];
            int allPartsIndex = 0;
            foreach (IReadOnlyList<string> pathParts in pathPartsList)
            {
                foreach (string part in pathParts)
                {
                    allParts[allPartsIndex++] = part;
                }
            }

            PartitionKind partitionKind = partitionKeyDefinition.Kind;

            this.Initialize();

            Span<byte> buffer = stackalloc byte[QueryPartitionProvider.InitialBufferSize];
            uint errorCode;
            uint serializedQueryExecutionInfoResultLength;

            string vectorEmbeddingPolicyString = vectorEmbeddingPolicy != null ?
                JsonConvert.SerializeObject(vectorEmbeddingPolicy) :
                null;

            unsafe
            {
                ServiceInteropWrapper.PartitionKeyRangesApiOptions partitionKeyRangesApiOptions =
                    new ServiceInteropWrapper.PartitionKeyRangesApiOptions()
                    {
                        bAllowDCount = Convert.ToInt32(allowDCount),
                        bAllowNonValueAggregateQuery = Convert.ToInt32(allowNonValueAggregateQuery),
                        bHasLogicalPartitionKey = Convert.ToInt32(hasLogicalPartitionKey),
                        bIsContinuationExpected = Convert.ToInt32(isContinuationExpected),
                        bRequireFormattableOrderByQuery = Convert.ToInt32(requireFormattableOrderByQuery),
                        bUseSystemPrefix = Convert.ToInt32(useSystemPrefix),
                        bHybridSearchSkipOrderByRewrite = Convert.ToInt32(hybridSearchSkipOrderByRewrite),
                        eGeospatialType = Convert.ToInt32(geospatialType),
                        ePartitionKind = Convert.ToInt32(partitionKind)
                    };

                fixed (byte* bytePtr = buffer)
                {
                    errorCode = ServiceInteropWrapper.GetPartitionKeyRangesFromQuery4(
                        this.serviceProvider,
                        querySpecJsonString,
                        partitionKeyRangesApiOptions,
                        allParts,
                        partsLengths,
                        (uint)partitionKeyDefinition.Paths.Count,
                        vectorEmbeddingPolicyString,
                        vectorEmbeddingPolicyString?.Length ?? 0,
                        new IntPtr(bytePtr),
                        (uint)buffer.Length,
                        out serializedQueryExecutionInfoResultLength);

                    if (errorCode == DISP_E_BUFFERTOOSMALL)
                    {
                        // Allocate on stack for smaller arrays, otherwise use heap.
                        buffer = serializedQueryExecutionInfoResultLength < 4096
                            ? stackalloc byte[(int)serializedQueryExecutionInfoResultLength]
                            : new byte[serializedQueryExecutionInfoResultLength];

                        fixed (byte* bytePtr2 = buffer)
                        {
                            errorCode = ServiceInteropWrapper.GetPartitionKeyRangesFromQuery4(
                                this.serviceProvider,
                                querySpecJsonString,
                                partitionKeyRangesApiOptions,
                                allParts,
                                partsLengths,
                                (uint)partitionKeyDefinition.Paths.Count,
                                vectorEmbeddingPolicyString,
                                vectorEmbeddingPolicyString?.Length ?? 0,
                                new IntPtr(bytePtr2),
                                (uint)buffer.Length,
                                out serializedQueryExecutionInfoResultLength);
                        }
                    }
                }
            }

            string serializedQueryExecutionInfo = Encoding.UTF8.GetString(buffer.Slice(0, (int)serializedQueryExecutionInfoResultLength));

            Exception exception = Marshal.GetExceptionForHR((int)errorCode);
            if (exception != null)
            {
                QueryPartitionProviderException queryPartitionProviderException;
                if (string.IsNullOrEmpty(serializedQueryExecutionInfo))
                {
                    queryPartitionProviderException = new UnexpectedQueryPartitionProviderException(
                        "Query service interop parsing hit an unexpected exception",
                        exception);
                }
                else
                {
                    queryPartitionProviderException = new ExpectedQueryPartitionProviderException(
                        serializedQueryExecutionInfo,
                        exception);
                }

                return TryCatch<PartitionedQueryExecutionInfoInternal>.FromException(
                    queryPartitionProviderException);
            }

            PartitionedQueryExecutionInfoInternal queryInfoInternal =
               JsonConvert.DeserializeObject<PartitionedQueryExecutionInfoInternal>(
                   serializedQueryExecutionInfo,
                   new JsonSerializerSettings
                   {
                       DateParseHandling = DateParseHandling.None,
                       MaxDepth = 64, // https://github.com/advisories/GHSA-5crp-9r3c-p9vr
                   });

            if (!this.ValidateQueryExecutionInfo(queryInfoInternal, out ArgumentException innerException))
            {
                return TryCatch<PartitionedQueryExecutionInfoInternal>.FromException(
                    new ExpectedQueryPartitionProviderException(
                        serializedQueryExecutionInfo,
                        innerException));
            }

            return TryCatch<PartitionedQueryExecutionInfoInternal>.FromResult(queryInfoInternal);
        }