Add tenant ID field to DDB object

Signed-off-by: Daniel Widdis <[email protected]>

plugin/src/main/java/org/opensearch/ml/sdkclient/DDBOpenSearchClient.java

@@ -116,6 +116,7 @@ public CompletionStage<PutDataObjectResponse> putDataObjectAsync(
         Boolean isMultiTenancyEnabled
     ) {
         final String id = request.id() != null ? request.id() : UUID.randomUUID().toString();
+        // This Default Tenant is only used on DynamoDB as it does not accept "null".
         final String tenantId = request.tenantId() != null ? request.tenantId() : DEFAULT_TENANT;
         final String tableName = request.index();
         final GetItemRequest getItemRequest = buildGetItemRequest(tenantId, id, request.index());
@@ -125,7 +126,12 @@ public CompletionStage<PutDataObjectResponse> putDataObjectAsync(
                 Long sequenceNumber = initOrIncrementSeqNo(getItemResponse);
                 String source = Strings.toString(MediaTypeRegistry.JSON, request.dataObject());
                 JsonNode jsonNode = OBJECT_MAPPER.readTree(source);
+                // This method is annotated @VisibleForTesting. The code should be made production or another method found
+                // This map is mutable, which we will take advantage of to add the tenantId
                 Map<String, AttributeValue> sourceMap = JsonTransformer.convertJsonObjectToDDBAttributeMap(jsonNode);
+                // Use of tenantId as key for DDB lookup doesn't propogate to OpenSearch document for search.
+                // Need to add it to the source map as well
+                sourceMap.put(TENANT_ID, AttributeValue.fromS(tenantId));
                 Map<String, AttributeValue> item = new HashMap<>();
                 item.put(TENANT_ID, AttributeValue.builder().s(tenantId).build());
                 item.put(RANGE_KEY, AttributeValue.builder().s(id).build());