[service-bus] Final work to expose AMQP body type encoding publicly (Azure#15295)

Finishing work to enable AMQP body type encoding

- We completely missed the schedule messages code path. Added, with tests.
- Consolidated the two files of AMQP messaging tests into the `public` branch so we also get min/max testing. Removed all references to non-published interfaces.

Author: richardpark-msft

sdk/servicebus/service-bus/CHANGELOG.md

@@ -1,6 +1,6 @@
 # Release History
 
-## 7.2.0-beta.1 (Unreleased)
+## 7.2.0-beta.1 (2021-05-18)
 
 ### New Features
 

sdk/servicebus/service-bus/review/service-bus.api.md

@@ -400,7 +400,7 @@ export interface ServiceBusMessageBatch {
     // @internal
     readonly _messageSpanContexts: SpanContext[];
     readonly sizeInBytes: number;
-    tryAddMessage(message: ServiceBusMessage, options?: TryAddOptions): boolean;
+    tryAddMessage(message: ServiceBusMessage | AmqpAnnotatedMessage, options?: TryAddOptions): boolean;
 }
 
 // @public
@@ -458,8 +458,8 @@ export interface ServiceBusSender {
     createMessageBatch(options?: CreateMessageBatchOptions): Promise<ServiceBusMessageBatch>;
     entityPath: string;
     isClosed: boolean;
-    scheduleMessages(messages: ServiceBusMessage | ServiceBusMessage[], scheduledEnqueueTimeUtc: Date, options?: OperationOptionsBase): Promise<Long[]>;
-    sendMessages(messages: ServiceBusMessage | ServiceBusMessage[] | ServiceBusMessageBatch, options?: OperationOptionsBase): Promise<void>;
+    scheduleMessages(messages: ServiceBusMessage | ServiceBusMessage[] | AmqpAnnotatedMessage | AmqpAnnotatedMessage[], scheduledEnqueueTimeUtc: Date, options?: OperationOptionsBase): Promise<Long[]>;
+    sendMessages(messages: ServiceBusMessage | ServiceBusMessage[] | ServiceBusMessageBatch | AmqpAnnotatedMessage | AmqpAnnotatedMessage[], options?: OperationOptionsBase): Promise<void>;
 }
 
 // @public

sdk/servicebus/service-bus/src/core/managementClient.ts

@@ -20,7 +20,8 @@ import {
   MessagingError,
   RequestResponseLink,
   SendRequestOptions,
-  RetryOptions
+  RetryOptions,
+  AmqpAnnotatedMessage
 } from "@azure/core-amqp";
 import { ConnectionContext } from "../connectionContext";
 import {
@@ -29,7 +30,9 @@ import {
   ServiceBusMessage,
   ServiceBusMessageImpl,
   toRheaMessage,
-  fromRheaMessage
+  fromRheaMessage,
+  updateScheduledTime,
+  updateMessageId
 } from "../serviceBusMessage";
 import { LinkEntity, RequestResponseLinkOptions } from "./linkEntity";
 import { managementClientLogger, receiverLogger, senderLogger, ServiceBusLogger } from "../log";
@@ -585,7 +588,7 @@ export class ManagementClient extends LinkEntity<RequestResponseLink> {
    */
   async scheduleMessages(
     scheduledEnqueueTimeUtc: Date,
-    messages: ServiceBusMessage[],
+    messages: ServiceBusMessage[] | AmqpAnnotatedMessage[],
     options?: OperationOptionsBase & SendManagementRequestOptions
   ): Promise<Long[]> {
     throwErrorIfConnectionClosed(this._context);
@@ -595,21 +598,28 @@ export class ManagementClient extends LinkEntity<RequestResponseLink> {
     const messageBody: any[] = [];
     for (let i = 0; i < messages.length; i++) {
       const item = messages[i];
-      if (!item.messageId) item.messageId = generate_uuid();
-      item.scheduledEnqueueTimeUtc = scheduledEnqueueTimeUtc;
 
       try {
-        const amqpMessage = toRheaMessage(item, defaultDataTransformer);
-
-        const entry: any = {
-          message: RheaMessageUtil.encode(amqpMessage),
-          "message-id": item.messageId
+        const rheaMessage = toRheaMessage(item, defaultDataTransformer);
+        updateMessageId(rheaMessage, rheaMessage.message_id || generate_uuid());
+        updateScheduledTime(rheaMessage, scheduledEnqueueTimeUtc);
+
+        const entry: {
+          message: Buffer;
+          ["message-id"]: ServiceBusMessage["messageId"];
+          ["partition-key"]?: ServiceBusMessage["partitionKey"];
+          [Constants.sessionIdMapKey]?: string | undefined;
+        } = {
+          message: RheaMessageUtil.encode(rheaMessage),
+          "message-id": rheaMessage.message_id
         };
-        if (item.sessionId) {
-          entry[Constants.sessionIdMapKey] = item.sessionId;
+
+        if (rheaMessage.group_id) {
+          entry[Constants.sessionIdMapKey] = rheaMessage.group_id;
         }
-        if (item.partitionKey) {
-          entry["partition-key"] = item.partitionKey;
+
+        if (rheaMessage.message_annotations?.[Constants.partitionKey]) {
+          entry["partition-key"] = rheaMessage.message_annotations[Constants.partitionKey];
         }
 
         // Will be required later for implementing Transactions
@@ -1319,3 +1329,41 @@ export class ManagementClient extends LinkEntity<RequestResponseLink> {
     }
   }
 }
+
+/**
+ * Converts an AmqpAnnotatedMessage or ServiceBusMessage into a properly formatted
+ * message for sending to the mgmt link for scheduling.
+ *
+ * @internal
+ * @hidden
+ */
+export function toScheduleableMessage(
+  item: ServiceBusMessage | AmqpAnnotatedMessage,
+  scheduledEnqueueTimeUtc: Date
+) {
+  const rheaMessage = toRheaMessage(item, defaultDataTransformer);
+  updateMessageId(rheaMessage, rheaMessage.message_id || generate_uuid());
+  updateScheduledTime(rheaMessage, scheduledEnqueueTimeUtc);
+
+  const entry: any = {
+    message: RheaMessageUtil.encode(rheaMessage),
+    "message-id": rheaMessage.message_id
+  };
+
+  rheaMessage.message_annotations = {
+    ...rheaMessage.message_annotations,
+    [Constants.scheduledEnqueueTime]: scheduledEnqueueTimeUtc
+  };
+
+  if (rheaMessage.group_id) {
+    entry[Constants.sessionIdMapKey] = rheaMessage.group_id;
+  }
+
+  const partitionKey =
+    rheaMessage.message_annotations && rheaMessage.message_annotations[Constants.partitionKey];
+
+  if (partitionKey) {
+    entry["partition-key"] = partitionKey;
+  }
+  return entry;
+}

sdk/servicebus/service-bus/src/sender.ts

@@ -51,7 +51,12 @@ export interface ServiceBusSender {
    * @throws `ServiceBusError` if the service returns an error while sending messages to the service.
    */
   sendMessages(
-    messages: ServiceBusMessage | ServiceBusMessage[] | ServiceBusMessageBatch,
+    messages:
+      | ServiceBusMessage
+      | ServiceBusMessage[]
+      | ServiceBusMessageBatch
+      | AmqpAnnotatedMessage
+      | AmqpAnnotatedMessage[],
     options?: OperationOptionsBase
   ): Promise<void>;
 
@@ -98,7 +103,11 @@ export interface ServiceBusSender {
    * @throws `ServiceBusError` if the service returns an error while scheduling messages.
    */
   scheduleMessages(
-    messages: ServiceBusMessage | ServiceBusMessage[],
+    messages:
+      | ServiceBusMessage
+      | ServiceBusMessage[]
+      | AmqpAnnotatedMessage
+      | AmqpAnnotatedMessage[],
     scheduledEnqueueTimeUtc: Date,
     options?: OperationOptionsBase
   ): Promise<Long[]>;
@@ -242,7 +251,11 @@ export class ServiceBusSenderImpl implements ServiceBusSender {
   }
 
   async scheduleMessages(
-    messages: ServiceBusMessage | ServiceBusMessage[],
+    messages:
+      | ServiceBusMessage
+      | ServiceBusMessage[]
+      | AmqpAnnotatedMessage
+      | AmqpAnnotatedMessage[],
     scheduledEnqueueTimeUtc: Date,
     options: OperationOptionsBase = {}
   ): Promise<Long[]> {
@@ -299,16 +312,14 @@ export class ServiceBusSenderImpl implements ServiceBusSender {
       ? sequenceNumbers
       : [sequenceNumbers];
     const cancelSchedulesMessagesOperationPromise = async (): Promise<void> => {
-      return this._context.getManagementClient(this._entityPath).cancelScheduledMessages(
-        sequenceNumbersToCancel,
-
-        {
+      return this._context
+        .getManagementClient(this._entityPath)
+        .cancelScheduledMessages(sequenceNumbersToCancel, {
           ...options,
           associatedLinkName: this._sender.name,
           requestName: "cancelScheduledMessages",
           timeoutInMs: this._retryOptions.timeoutInMs
-        }
-      );
+        });
     };
     const config: RetryConfig<void> = {
       operation: cancelSchedulesMessagesOperationPromise,