Merge branch 'Azure:main' into main

.github/CODEOWNERS

@@ -168,6 +168,9 @@
 # PRLabel: %Cognitive - Language
 /sdk/cognitivelanguage/ @minhanh-phan @deyaaeldeen
 
+# PRLabel: %OpenAI
+/sdk/openai/ @deyaaeldeen @minhanh-phan
+
 # PRLabel: %Schema Registry
 /sdk/schemaregistry/ @deyaaeldeen
 

eng/emitter-package.json

@@ -1,7 +1,6 @@
 {
   "main": "dist/src/index.js",
   "dependencies": {
-    "@azure-tools/typespec-ts": "0.13.0",
-    "@typespec/versioning": "0.44.0"
+    "@azure-tools/typespec-ts": "latest"
   }
 }

sdk/cosmosdb/cosmos/review/cosmos.api.md

@@ -32,6 +32,11 @@ export interface Agent {
 // @public (undocumented)
 export type AggregateType = "Average" | "Count" | "Max" | "Min" | "Sum";
 
+// @public (undocumented)
+export type BulkOperationResponse = OperationResponse[] & {
+    diagnostics: CosmosDiagnostics;
+};
+
 // @public (undocumented)
 export const BulkOperationType: {
     readonly Create: "Create";
@@ -175,7 +180,7 @@ export class ClientContext {
         diagnosticContext?: CosmosDiagnosticContext;
     }): Promise<Response_2<T & Resource>>;
     // (undocumented)
-    queryPartitionKeyRanges(collectionLink: string, query?: string | SqlQuerySpec, options?: FeedOptions): QueryIterator<PartitionKeyRange>;
+    queryPartitionKeyRanges(collectionLink: string, diagnosticContext: CosmosDiagnosticContext, query?: string | SqlQuerySpec, options?: FeedOptions): QueryIterator<PartitionKeyRange>;
     // (undocumented)
     read<T>({ path, resourceType, resourceId, options, partitionKey, diagnosticContext, }: {
         path: string;
@@ -814,14 +819,16 @@ export interface FeedOptions extends SharedOptions {
 
 // @public (undocumented)
 export class FeedResponse<TResource> {
-    constructor(resources: TResource[], headers: CosmosHeaders, hasMoreResults: boolean);
+    constructor(resources: TResource[], headers: CosmosHeaders, hasMoreResults: boolean, diagnostics: CosmosDiagnostics);
     // (undocumented)
     get activityId(): string;
     // (undocumented)
     get continuation(): string;
     // (undocumented)
     get continuationToken(): string;
     // (undocumented)
+    readonly diagnostics: CosmosDiagnostics;
+    // (undocumented)
     readonly hasMoreResults: boolean;
     // (undocumented)
     get indexMetrics(): string;
@@ -958,7 +965,7 @@ export class ItemResponse<T extends ItemDefinition> extends ResourceResponse<T &
 export class Items {
     constructor(container: Container, clientContext: ClientContext);
     batch(operations: OperationInput[], partitionKey?: string, options?: RequestOptions): Promise<Response_2<OperationResponse[]>>;
-    bulk(operations: OperationInput[], bulkOptions?: BulkOptions, options?: RequestOptions): Promise<OperationResponse[]>;
+    bulk(operations: OperationInput[], bulkOptions?: BulkOptions, options?: RequestOptions): Promise<BulkOperationResponse>;
     changeFeed(partitionKey: string | number | boolean, changeFeedOptions?: ChangeFeedOptions): ChangeFeedIterator<any>;
     changeFeed(changeFeedOptions?: ChangeFeedOptions): ChangeFeedIterator<any>;
     changeFeed<T>(partitionKey: string | number | boolean, changeFeedOptions?: ChangeFeedOptions): ChangeFeedIterator<T>;
@@ -1333,7 +1340,7 @@ export interface QueryInfo {
 // @public
 export class QueryIterator<T> {
     // Warning: (ae-forgotten-export) The symbol "FetchFunctionCallback" needs to be exported by the entry point index.d.ts
-    constructor(clientContext: ClientContext, query: SqlQuerySpec | string, options: FeedOptions, fetchFunctions: FetchFunctionCallback | FetchFunctionCallback[], resourceLink?: string, resourceType?: ResourceType);
+    constructor(clientContext: ClientContext, query: SqlQuerySpec | string, options: FeedOptions, fetchFunctions: FetchFunctionCallback | FetchFunctionCallback[], diagnosticContext?: CosmosDiagnosticContext, resourceLink?: string, resourceType?: ResourceType);
     fetchAll(): Promise<FeedResponse<T>>;
     fetchNext(): Promise<FeedResponse<T>>;
     getAsyncIterator(): AsyncIterable<FeedResponse<T>>;
@@ -1588,11 +1595,11 @@ export interface Resource {
 
 // @public (undocumented)
 export class ResourceResponse<TResource> {
-    constructor(resource: TResource | undefined, headers: CosmosHeaders_2, statusCode: StatusCode, diagnostics?: CosmosDiagnostics, substatus?: SubStatusCode);
+    constructor(resource: TResource | undefined, headers: CosmosHeaders_2, statusCode: StatusCode, diagnostics: CosmosDiagnostics, substatus?: SubStatusCode);
     // (undocumented)
     get activityId(): string;
     // (undocumented)
-    readonly diagnostics?: CosmosDiagnostics;
+    readonly diagnostics: CosmosDiagnostics;
     // (undocumented)
     get etag(): string;
     // (undocumented)
@@ -1642,7 +1649,7 @@ interface Response_2<T> {
     // (undocumented)
     code?: number;
     // (undocumented)
-    diagnostics?: CosmosDiagnostics;
+    diagnostics: CosmosDiagnostics;
     // (undocumented)
     headers: CosmosHeaders;
     // (undocumented)
@@ -2202,7 +2209,7 @@ export class Users {
 
 // Warnings were encountered during analysis:
 //
-// src/ClientContext.ts:90:5 - (ae-forgotten-export) The symbol "CosmosDiagnosticContext" needs to be exported by the entry point index.d.ts
+// src/ClientContext.ts:91:5 - (ae-forgotten-export) The symbol "CosmosDiagnosticContext" needs to be exported by the entry point index.d.ts
 
 // (No @packageDocumentation comment for this package)
 

sdk/cosmosdb/cosmos/src/ClientContext.ts

@@ -31,6 +31,7 @@ import { BulkOptions } from "./utils/batch";
 import { sanitizeEndpoint } from "./utils/checkURL";
 import { AzureLogger, createClientLogger } from "@azure/logger";
 import { CosmosDiagnosticContext } from "./CosmosDiagnosticsContext";
+import { MetadataLookUpType } from "./CosmosDiagnostics";
 
 const logger: AzureLogger = createClientLogger("ClientContext");
 
@@ -229,22 +230,31 @@ export class ClientContext {
 
   public queryPartitionKeyRanges(
     collectionLink: string,
+    diagnosticContext: CosmosDiagnosticContext,
     query?: string | SqlQuerySpec,
     options?: FeedOptions
   ): QueryIterator<PartitionKeyRange> {
     const path = getPathFromLink(collectionLink, ResourceType.pkranges);
     const id = getIdFromLink(collectionLink);
-    const cb: FetchFunctionCallback = (innerOptions) => {
-      return this.queryFeed({
+    const cb: FetchFunctionCallback = async (
+      innerOptions,
+      diagnosticCtx: CosmosDiagnosticContext
+    ) => {
+      const response = await this.queryFeed({
         path,
         resourceType: ResourceType.pkranges,
         resourceId: id,
         resultFn: (result) => result.PartitionKeyRanges,
         query,
         options: innerOptions,
       });
+      diagnosticCtx.recordMetaDataLookup(
+        response.diagnostics,
+        MetadataLookUpType.PartitionKeyRangeLookUp
+      );
+      return response;
     };
-    return new QueryIterator<PartitionKeyRange>(this, query, options, cb);
+    return new QueryIterator<PartitionKeyRange>(this, query, options, cb, diagnosticContext);
   }
 
   public async delete<T>({

sdk/cosmosdb/cosmos/src/CosmosDiagnosticsContext.ts

@@ -95,6 +95,25 @@ export class CosmosDiagnosticContext {
     this.locationEndpointsContacted.set(location.databaseAccountEndpoint, location);
   }
 
+  public reset(): void {
+    this.requestStartTimeUTCinMs = getCurrentTimestampInMs();
+    this.requestEndTimeUTCinMs = getCurrentTimestampInMs();
+    this.retryStartTimeUTCinMs = 0;
+    this.headers = {};
+    this.retryAttemptNumber = 0;
+    this.failedAttempts = [];
+    this.metadataLookups = [];
+    this.requestPayloadLength = 0;
+    this.responsePayloadLength = 0;
+    this.locationEndpointsContacted = new Map();
+  }
+
+  public resetAndGetDiagnostics(): CosmosDiagnostics {
+    const diagnostic = this.getDiagnostics();
+    this.reset();
+    return diagnostic;
+  }
+
   public getDiagnostics(): CosmosDiagnostics {
     this.recordSessionEnd();
     return new CosmosDiagnostics(

sdk/cosmosdb/cosmos/src/client/ClientUtils.ts

@@ -0,0 +1,19 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+import { MetadataLookUpType } from "../CosmosDiagnostics";
+import { CosmosDiagnosticContext } from "../CosmosDiagnosticsContext";
+import { PartitionKeyDefinition } from "../documents";
+import { Container } from "./Container";
+
+export async function readAndRecordPartitionKeyDefinition(container: Container): Promise<{
+  diagnosticContext: CosmosDiagnosticContext;
+  partitionKeyDefinition: PartitionKeyDefinition;
+}> {
+  const diagnosticContext: CosmosDiagnosticContext = new CosmosDiagnosticContext();
+
+  const { resource: partitionKeyDefinition, diagnostics } =
+    await container.readPartitionKeyDefinition();
+  diagnosticContext.recordMetaDataLookup(diagnostics, MetadataLookUpType.PartitionKeyRangeLookUp);
+  return { diagnosticContext, partitionKeyDefinition };
+}

sdk/cosmosdb/cosmos/src/client/Container/Container.ts

@@ -25,6 +25,7 @@ import { Offer, OfferDefinition } from "../Offer";
 import { OfferResponse } from "../Offer/OfferResponse";
 import { Resource } from "../Resource";
 import { getEmptyCosmosDiagnostics } from "../../CosmosDiagnostics";
+import { CosmosDiagnosticContext } from "../../CosmosDiagnosticsContext";
 
 /**
  * Operations for reading, replacing, or deleting a specific, existing container by id.
@@ -266,7 +267,12 @@ export class Container {
 
   public readPartitionKeyRanges(feedOptions?: FeedOptions): QueryIterator<PartitionKeyRange> {
     feedOptions = feedOptions || {};
-    return this.clientContext.queryPartitionKeyRanges(this.url, undefined, feedOptions);
+    return this.clientContext.queryPartitionKeyRanges(
+      this.url,
+      new CosmosDiagnosticContext(),
+      undefined,
+      feedOptions
+    );
   }
 
   /**

sdk/cosmosdb/cosmos/src/client/Item/Item.ts

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 import { ClientContext } from "../../ClientContext";
+import { CosmosDiagnosticContext } from "../../CosmosDiagnosticsContext";
 import {
   createDocumentUri,
   getIdFromLink,
@@ -13,6 +14,7 @@ import { PartitionKey } from "../../documents";
 import { extractPartitionKey, undefinedPartitionKey } from "../../extractPartitionKey";
 import { RequestOptions, Response } from "../../request";
 import { PatchRequestBody } from "../../utils/patch";
+import { readAndRecordPartitionKeyDefinition } from "../ClientUtils";
 import { Container } from "../Container";
 import { Resource } from "../Resource";
 import { ItemDefinition } from "./ItemDefinition";
@@ -74,10 +76,11 @@ export class Item {
   public async read<T extends ItemDefinition = any>(
     options: RequestOptions = {}
   ): Promise<ItemResponse<T>> {
+    let diagnosticContext: CosmosDiagnosticContext;
     if (this.partitionKey === undefined) {
-      const { resource: partitionKeyDefinition } =
-        await this.container.readPartitionKeyDefinition();
-      this.partitionKey = undefinedPartitionKey(partitionKeyDefinition);
+      const partitionKeyResponse = await readAndRecordPartitionKeyDefinition(this.container);
+      this.partitionKey = undefinedPartitionKey(partitionKeyResponse.partitionKeyDefinition);
+      diagnosticContext = partitionKeyResponse.diagnosticContext;
     }
 
     const path = getPathFromLink(this.url);
@@ -90,6 +93,7 @@ export class Item {
         resourceId: id,
         options,
         partitionKey: this.partitionKey,
+        diagnosticContext,
       });
     } catch (error: any) {
       if (error.code !== StatusCodes.NotFound) {
@@ -139,10 +143,11 @@ export class Item {
     body: T,
     options: RequestOptions = {}
   ): Promise<ItemResponse<T>> {
+    let diagnosticContext: CosmosDiagnosticContext;
     if (this.partitionKey === undefined) {
-      const { resource: partitionKeyDefinition } =
-        await this.container.readPartitionKeyDefinition();
-      this.partitionKey = extractPartitionKey(body, partitionKeyDefinition);
+      const partitionKeyResponse = await readAndRecordPartitionKeyDefinition(this.container);
+      this.partitionKey = extractPartitionKey(body, partitionKeyResponse.partitionKeyDefinition);
+      diagnosticContext = partitionKeyResponse.diagnosticContext;
     }
 
     const err = {};
@@ -160,6 +165,7 @@ export class Item {
       resourceId: id,
       options,
       partitionKey: this.partitionKey,
+      diagnosticContext,
     });
     return new ItemResponse(
       response.result,
@@ -182,10 +188,11 @@ export class Item {
   public async delete<T extends ItemDefinition = any>(
     options: RequestOptions = {}
   ): Promise<ItemResponse<T>> {
+    let diagnosticContext: CosmosDiagnosticContext;
     if (this.partitionKey === undefined) {
-      const { resource: partitionKeyDefinition } =
-        await this.container.readPartitionKeyDefinition();
-      this.partitionKey = undefinedPartitionKey(partitionKeyDefinition);
+      const partitionKeyResponse = await readAndRecordPartitionKeyDefinition(this.container);
+      this.partitionKey = undefinedPartitionKey(partitionKeyResponse.partitionKeyDefinition);
+      diagnosticContext = partitionKeyResponse.diagnosticContext;
     }
 
     const path = getPathFromLink(this.url);
@@ -197,6 +204,7 @@ export class Item {
       resourceId: id,
       options,
       partitionKey: this.partitionKey,
+      diagnosticContext,
     });
     return new ItemResponse(
       response.result,
@@ -220,10 +228,11 @@ export class Item {
     body: PatchRequestBody,
     options: RequestOptions = {}
   ): Promise<ItemResponse<T>> {
+    let diagnosticContext: CosmosDiagnosticContext;
     if (this.partitionKey === undefined) {
-      const { resource: partitionKeyDefinition } =
-        await this.container.readPartitionKeyDefinition();
-      this.partitionKey = extractPartitionKey(body, partitionKeyDefinition);
+      const partitionKeyResponse = await readAndRecordPartitionKeyDefinition(this.container);
+      this.partitionKey = extractPartitionKey(body, partitionKeyResponse.partitionKeyDefinition);
+      diagnosticContext = partitionKeyResponse.diagnosticContext;
     }
 
     const path = getPathFromLink(this.url);
@@ -236,6 +245,7 @@ export class Item {
       resourceId: id,
       options,
       partitionKey: this.partitionKey,
+      diagnosticContext,
     });
     return new ItemResponse(
       response.result,