From a6d6c9eaff3848c19f7a9dcc111741d94490a9f4 Mon Sep 17 00:00:00 2001
From: tywalch <tywalch@gmail.com>
Date: Thu, 28 Mar 2024 16:04:18 -0400
Subject: [PATCH] Fixes 362 This branch fixes issue #362. The root cause of the
 issue was that one of the composite attributes was hidden. Because the cursor
 was created using the "formatted" return item, it was the hidden composite
 was not available to the cursor formatting function. We now check for this
 unique case and perform some additional calculations. Ideally this reduces
 the surface area of performance implications to only entities that have a
 hidden composite attribute defined.

---
 src/entity.js | 55 ++++++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 46 insertions(+), 9 deletions(-)

diff --git a/src/entity.js b/src/entity.js
index e46a78cd..25fd5064 100644
--- a/src/entity.js
+++ b/src/entity.js
@@ -733,6 +733,9 @@ class Entity {
     let count = 0;
     let hydratedUnprocessed = [];
     const shouldHydrate = config.hydrate && method === MethodTypes.query;
+    let allRawItems = [];
+    let indexHasHiddenComposites = (this.model.lookup.indexHasHiddenComposites[indexName] || this.model.lookup.indexHasHiddenComposites[TableIndex]);
+    let hasCountOption = typeof config.count === "number";
     do {
       let limit = max === undefined ? parameters.Limit : max - count;
       let response = await this._exec(
@@ -747,6 +750,7 @@ class Entity {
         ...config,
         includeKeys: shouldHydrate || config.includeKeys,
         ignoreOwnership: shouldHydrate || config.ignoreOwnership,
+        includeRaw: hasCountOption && indexHasHiddenComposites,
       });
 
       if (config.raw) {
@@ -774,13 +778,16 @@ class Entity {
         }
       } else if (Array.isArray(response.data)) {
         let prevCount = count
-        if (!!max || !!config.count) {
+        if (!!max || hasCountOption) {
           count += response.data.length;
         }
         let items = response.data;
-        const moreItemsThanRequired = !!config.count && count > config.count;
+        const moreItemsThanRequired = hasCountOption && count > config.count;
+        let rawItems = response.raw || [];
         if (moreItemsThanRequired) {
-          items = items.slice(0, config.count - prevCount);
+          const endIndex = config.count - prevCount;
+          items = items.slice(0, endIndex);
+          rawItems = rawItems.slice(0, endIndex);
         }
         if (shouldHydrate) {
           const hydrated = await this.hydrate(
@@ -794,9 +801,10 @@ class Entity {
           );
         }
         results = [...results, ...items];
+        allRawItems = [...allRawItems, ...rawItems];
         if (moreItemsThanRequired || count === config.count) {
-          const lastItem = results[results.length - 1];
-          ExclusiveStartKey = this._fromCompositeToKeysByIndex({ indexName, provided: lastItem });
+          const lastRawItem = rawItems[rawItems.length - 1];
+          ExclusiveStartKey = this._trimKeysToIndex({indexName, provided: lastRawItem});
           break;
         }
       } else {
@@ -955,6 +963,7 @@ class Entity {
       stackTrace = new e.ElectroError(e.ErrorCodes.AWSError);
     }
     try {
+      let raw = {};
       let results = {};
       if (validations.isFunction(config.parse)) {
         results = config.parse(config, response);
@@ -980,18 +989,24 @@ class Entity {
               this.ownsKeys(response.Item)) ||
             this.ownsItem(response.Item)
           ) {
+            if (config.includeRaw) {
+              raw = results.Item;
+            }
             results = this.model.schema.formatItemForRetrieval(
               response.Item,
               config,
             );
             if (Object.keys(results).length === 0) {
+              raw = null;
               results = null;
             }
           } else if (!config._objectOnEmpty) {
+            raw = null;
             results = null;
           }
         } else if (response.Items) {
           results = [];
+          raw = [];
           for (let item of response.Items) {
             if (
               (config.ignoreOwnership &&
@@ -1008,6 +1023,9 @@ class Entity {
               );
               if (Object.keys(record).length > 0) {
                 results.push(record);
+                if (config.includeRaw) {
+                  raw.push(item);
+                }
               }
             }
           }
@@ -1017,6 +1035,7 @@ class Entity {
             config,
           );
           if (Object.keys(results).length === 0) {
+            raw = null;
             results = null;
           }
         } else if (config._objectOnEmpty) {
@@ -1027,18 +1046,28 @@ class Entity {
           };
         } else {
           results = null;
+          raw = null;
         }
       }
 
+      let formatted = {
+        data: results,
+      }
+
+      if (config.includeRaw) {
+        formatted.raw = raw;
+      }
+
       if (config._isPagination || response.LastEvaluatedKey) {
         const nextPage = this._formatReturnPager(
           config,
           response.LastEvaluatedKey,
         );
-        return { cursor: nextPage || null, data: results };
+
+        formatted.cursor = nextPage || null;
       }
 
-      return { data: results };
+      return formatted;
     } catch (err) {
       if (
         config.originalErr ||
@@ -4044,7 +4073,7 @@ class Entity {
         }
       }
 
-      let definition= {
+      let definition = {
         pk,
         sk,
         hasSk,
@@ -4498,7 +4527,14 @@ class Entity {
       getClient: () => this.client,
       isRoot: true,
     });
-
+    const indexHasHiddenComposites = {};
+    for (let hiddenAttr of schema.hiddenAttributes.values()) {
+      if (facets.byAttr[hiddenAttr]) {
+        for (const facet of facets.byAttr[hiddenAttr]) {
+          indexHasHiddenComposites[facet.index] = true;
+        }
+      }
+    }
     let filters = this._normalizeFilters(model.filters);
     // todo: consider a rename
     let prefixes = this._normalizeKeyFixings({
@@ -4562,6 +4598,7 @@ class Entity {
         clusteredIndexes,
         indexHasSortKeys,
         indexHasSubCollections,
+        indexHasHiddenComposites,
       },
       translations: {
         keys: indexField,