Support scanning GSI

package.json

@@ -26,7 +26,7 @@
     "integration-tests": "npx vitest run --dir test/integration --config ./test/integration/vitest.config.ts",
     "all-checks": "npm run local-checks && npm run integration-tests",
     "deploy": "cd test/examples && aws cloudformation deploy --template-file template.yaml --stack-name \"${STACK_NAME-entity-store-test-stack}\" --no-fail-on-empty-changeset",
-    "deploy-and-all-checks": "npm run deploy && npm run integration-tests",
+    "deploy-and-all-checks": "npm run deploy && npm run all-checks",
     "build": "rm -rf dist && tsc -p tsconfig-build-esm.json && tsc -p tsconfig-build-cjs.json && echo '{\"type\": \"module\"}' > dist/esm/package.json && echo '{\"type\": \"commonjs\"}' > dist/cjs/package.json",
     "prepublishOnly": "npm run local-checks && npm run build",
     "check-examples": "cd examples && npm run local-checks",

src/lib/internal/common/gsiQueryCommon.ts

@@ -1,7 +1,7 @@
 import { GsiGenerators } from '../../entities'
 import { EntityContext } from '../entityContext'
 import { throwError } from '../../util'
-import { AdvancedGsiQueryOnePageOptions } from '../../singleEntityAdvancedOperations'
+import { WithGsiId } from '../../singleEntityOperations'
 
 export interface GsiDetails {
   id: string
@@ -13,9 +13,9 @@ export interface GsiDetails {
 // TODO - needs more testing
 export function findGsiDetails<TItem extends TPKSource & TSKSource, TPKSource, TSKSource>(
   entityContext: EntityContext<TItem, TPKSource, TSKSource>,
-  options: AdvancedGsiQueryOnePageOptions
+  withGsiId: WithGsiId
 ): GsiDetails {
-  const entityGsi = findEntityGsi(entityContext, options)
+  const entityGsi = findEntityGsi(entityContext, withGsiId)
   const tableGsi = findGsiTableDetails(entityContext, entityGsi.gsiId)
   return {
     id: entityGsi.gsiId,
@@ -27,7 +27,7 @@ export function findGsiDetails<TItem extends TPKSource & TSKSource, TPKSource, T
 
 function findEntityGsi<TItem extends TPKSource & TSKSource, TPKSource, TSKSource>(
   { entity: { type, gsis } }: EntityContext<TItem, TPKSource, TSKSource>,
-  options: AdvancedGsiQueryOnePageOptions
+  withGsiId: WithGsiId
 ): { gsiId: string; gsiGenerators: GsiGenerators } {
   const entityGsiCount = Object.keys(gsis ?? {}).length
   if (!gsis || entityGsiCount === 0)
@@ -39,13 +39,13 @@ function findEntityGsi<TItem extends TPKSource & TSKSource, TPKSource, TSKSource
       gsiGenerators: onlyGsi[1]
     }
   }
-  if (!options.gsiId)
+  if (!withGsiId.gsiId)
     throw new Error(
       `Entity type ${type} has multiple GSIs but no GSI ID (.gsiId) was specified on query options`
     )
   return {
-    gsiId: options.gsiId,
-    gsiGenerators: gsis[options.gsiId]
+    gsiId: withGsiId.gsiId,
+    gsiGenerators: gsis[withGsiId.gsiId]
   }
 }
 

src/lib/internal/common/queryAndScanCommon.ts

@@ -15,6 +15,7 @@ import {
   AdvancedScanOnePageOptions,
   ConsumedCapacitiesMetadata
 } from '../../singleEntityAdvancedOperations'
+import { GsiDetails } from './gsiQueryCommon'
 
 export interface QueryScanOperationConfiguration<
   TCommandInput extends ScanCommandInput & QueryCommandInput,
@@ -29,10 +30,16 @@ export interface QueryScanOperationConfiguration<
 export function configureScanOperation(
   { dynamoDB, tableName }: Pick<EntityContext<never, never, never>, 'tableName' | 'dynamoDB'>,
   options: AdvancedScanOnePageOptions,
-  allPages: boolean
+  allPages: boolean,
+  gsiDetails?: GsiDetails
 ): QueryScanOperationConfiguration<ScanCommandInput, ScanCommandOutput> {
   return {
-    ...configureOperation(tableName, options, allPages, undefined),
+    ...configureOperation(
+      tableName,
+      options,
+      allPages,
+      gsiDetails ? { IndexName: gsiDetails.tableIndexName } : undefined
+    ),
     allPageOperation: dynamoDB.scanAllPages,
     onePageOperation: dynamoDB.scanOnePage
   }
@@ -55,7 +62,7 @@ export function configureOperation(
   tableName: string,
   options: AdvancedQueryOnePageOptions,
   allPages: boolean,
-  queryParamsParts?: Omit<QueryCommandInput, 'TableName' | 'ExclusiveStartKey' | 'Limit'>
+  paramsParts?: Omit<QueryCommandInput, 'TableName' | 'ExclusiveStartKey' | 'Limit'>
 ): { operationParams: ScanCommandInput & QueryCommandInput; useAllPageOperation: boolean } {
   const { limit, exclusiveStartKey, consistentRead } = options
   return {
@@ -64,7 +71,7 @@ export function configureOperation(
       ...(limit ? { Limit: limit } : {}),
       ...(exclusiveStartKey ? { ExclusiveStartKey: exclusiveStartKey } : {}),
       ...(consistentRead !== undefined ? { ConsistentRead: consistentRead } : {}),
-      ...(queryParamsParts ? queryParamsParts : {}),
+      ...(paramsParts ? paramsParts : {}),
       ...returnConsumedCapacityParam(options)
     },
     useAllPageOperation: allPages

src/lib/internal/singleEntity/scanItems.ts

@@ -5,13 +5,15 @@ import {
   parseResultsForEntity
 } from '../common/queryAndScanCommon'
 import { AdvancedCollectionResponse, AdvancedScanOnePageOptions } from '../../singleEntityAdvancedOperations'
+import { GsiDetails } from '../common/gsiQueryCommon'
 
 export async function scanItems<TItem extends TPKSource & TSKSource, TPKSource, TSKSource>(
   context: EntityContext<TItem, TPKSource, TSKSource>,
   options: AdvancedScanOnePageOptions,
-  allPages: boolean
+  allPages: boolean,
+  gsiDetails?: GsiDetails
 ): Promise<AdvancedCollectionResponse<TItem>> {
-  const scanConfig = configureScanOperation(context, options, allPages)
+  const scanConfig = configureScanOperation(context, options, allPages, gsiDetails)
   const result = await executeQueryOrScan(scanConfig, context.logger, context.entity.type)
   return parseResultsForEntity(context, result)
 }

src/lib/internal/singleEntity/tableBackedSingleEntityAdvancedOperations.ts

@@ -9,6 +9,8 @@ import {
   AdvancedGetResponse,
   AdvancedGsiQueryAllOptions,
   AdvancedGsiQueryOnePageOptions,
+  AdvancedGsiScanAllOptions,
+  AdvancedGsiScanOnePageOptions,
   AdvancedPutOptions,
   AdvancedPutResponse,
   AdvancedQueryAllOptions,
@@ -51,6 +53,10 @@ export function tableBackedSingleEntityAdvancedOperations<
   entity: Entity<TItem, TPKSource, TSKSource>,
   entityContext: EntityContext<TItem, TPKSource, TSKSource>
 ): SingleEntityAdvancedOperations<TItem, TPKSource, TSKSource> {
+  function checkAllowScans() {
+    if (!table.allowScans) throw new Error('Scan operations are disabled for this store')
+  }
+
   return {
     async put(item: TItem, options?: AdvancedPutOptions): Promise<AdvancedPutResponse> {
       return await putItem(entityContext, item, options)
@@ -170,15 +176,25 @@ export function tableBackedSingleEntityAdvancedOperations<
     },
 
     async scanAll(options: AdvancedScanAllOptions = {}) {
-      if (!table.allowScans) throw new Error('Scan operations are disabled for this store')
+      checkAllowScans()
       return await scanItems(entityContext, options, true)
     },
 
     async scanOnePage(options: AdvancedScanOnePageOptions = {}) {
-      if (!table.allowScans) throw new Error('Scan operations are disabled for this store')
+      checkAllowScans()
       return await scanItems(entityContext, options, false)
     },
 
+    async scanAllWithGsi(options: AdvancedGsiScanAllOptions = {}) {
+      checkAllowScans()
+      return await scanItems(entityContext, options, true, findGsiDetails(entityContext, options))
+    },
+
+    async scanOnePageWithGsi(options: AdvancedGsiScanOnePageOptions = {}) {
+      checkAllowScans()
+      return await scanItems(entityContext, options, false, findGsiDetails(entityContext, options))
+    },
+
     async batchPut(items: TItem[], options?: BatchPutOptions): Promise<AdvancedBatchWriteResponse> {
       return await batchPutItems(entityContext, items, options)
     },

src/lib/internal/singleEntity/tableBackedSingleEntityOperations.ts

@@ -5,6 +5,8 @@ import {
   GetOptions,
   GsiQueryAllOptions,
   GsiQueryOnePageOptions,
+  GsiScanAllOptions,
+  GsiScanOnePageOptions,
   OnePageResponse,
   PutOptions,
   QueryAllOptions,
@@ -123,6 +125,14 @@ export function tableBackedSingleEntityOperations<TItem extends TPKSource & TSKS
 
     async scanOnePage(options: ScanOnePageOptions = {}) {
       return await advancedOperations.scanOnePage(options)
+    },
+
+    async scanAllWithGsi(options: GsiScanAllOptions = {}) {
+      return (await advancedOperations.scanAllWithGsi(options)).items
+    },
+
+    async scanOnePageWithGsi(options: GsiScanOnePageOptions = {}) {
+      return await advancedOperations.scanOnePageWithGsi(options)
     }
   }
 }

src/lib/singleEntityAdvancedOperations.ts

@@ -3,6 +3,8 @@ import {
   GetOptions,
   GsiQueryAllOptions,
   GsiQueryOnePageOptions,
+  GsiScanAllOptions,
+  GsiScanOnePageOptions,
   OnePageResponse,
   PutOptions,
   QueryAllOptions,
@@ -94,6 +96,10 @@ export interface SingleEntityAdvancedOperations<TItem extends TPKSource & TSKSou
 
   scanOnePage(options?: AdvancedScanOnePageOptions): Promise<AdvancedCollectionResponse<TItem>>
 
+  scanAllWithGsi(options?: AdvancedGsiScanAllOptions): Promise<AdvancedCollectionResponse<TItem>>
+
+  scanOnePageWithGsi(options?: AdvancedGsiScanOnePageOptions): Promise<AdvancedCollectionResponse<TItem>>
+
   batchPut(item: TItem[], options?: BatchPutOptions): Promise<AdvancedBatchWriteResponse>
 
   batchDelete<TKeySource extends TPKSource & TSKSource>(
@@ -136,6 +142,10 @@ export type AdvancedScanAllOptions = ScanAllOptions & ReturnConsumedCapacityOpti
 
 export type AdvancedScanOnePageOptions = ScanOnePageOptions & ReturnConsumedCapacityOption
 
+export type AdvancedGsiScanAllOptions = GsiScanAllOptions & ReturnConsumedCapacityOption
+
+export type AdvancedGsiScanOnePageOptions = GsiScanOnePageOptions & ReturnConsumedCapacityOption
+
 export type BatchPutOptions = BatchOptions &
   ReturnConsumedCapacityOption &
   ReturnItemCollectionMetricsOption & {

src/lib/singleEntityOperations.ts

@@ -70,6 +70,10 @@ export interface SingleEntityOperations<TItem extends TPKSource & TSKSource, TPK
   scanAll(options?: ScanAllOptions): Promise<TItem[]>
 
   scanOnePage(options?: ScanOnePageOptions): Promise<OnePageResponse<TItem>>
+
+  scanAllWithGsi(options?: GsiScanAllOptions): Promise<TItem[]>
+
+  scanOnePageWithGsi(options?: GsiScanOnePageOptions): Promise<OnePageResponse<TItem>>
 }
 
 export interface PutOptions {
@@ -240,21 +244,17 @@ export interface QueryOnePageOptions {
   consistentRead?: boolean
 }
 
-export interface GsiQueryAllOptions extends Omit<QueryAllOptions, 'consistentRead'> {
+export interface WithGsiId {
   /**
    * If an entity has multiple GSIs then this property must be used to specify which GSI to use
    * @default use the only GSI on the entity
    */
   gsiId?: string
 }
 
-export interface GsiQueryOnePageOptions extends Omit<QueryOnePageOptions, 'consistentRead'> {
-  /**
-   * If an entity has multiple GSIs then this property must be used to specify which GSI to use
-   * @default use the only GSI on the entity
-   */
-  gsiId?: string
-}
+export type GsiQueryAllOptions = Omit<QueryAllOptions, 'consistentRead'> & WithGsiId
+
+export type GsiQueryOnePageOptions = Omit<QueryOnePageOptions, 'consistentRead'> & WithGsiId
 
 export interface SkQueryRange {
   /**
@@ -309,6 +309,10 @@ export interface ScanOnePageOptions {
   consistentRead?: boolean
 }
 
+export type GsiScanAllOptions = WithGsiId
+
+export type GsiScanOnePageOptions = Omit<ScanOnePageOptions, 'consistentRead'> & WithGsiId
+
 export interface OnePageResponse<TItem> {
   /**
    * Result of query

test/integration/singleTable/chickens.test.ts

@@ -36,6 +36,38 @@ describe('chickens', () => {
     ).toBeUndefined()
   })
 
+  describe('scan', () => {
+    test('scan GSI disabled by default', async () => {
+      const chickenStore = (await initialize()).for(CHICKEN_ENTITY)
+      expect(async () => await chickenStore.scanOnePageWithGsi()).rejects.toThrowError(
+        'Scan operations are disabled for this store'
+      )
+      expect(async () => await chickenStore.scanAllWithGsi()).rejects.toThrowError(
+        'Scan operations are disabled for this store'
+      )
+    })
+
+    test('scan GSI', async () => {
+      const chickenStore = (await initialize({ allowScans: true })).for(CHICKEN_ENTITY)
+      await chickenStore.put(ginger)
+      await chickenStore.put(babs)
+      await chickenStore.put(bunty)
+
+      const scanAllGsiResult = await chickenStore.scanAllWithGsi()
+      expect(scanAllGsiResult).toEqual([ginger, babs, bunty])
+
+      const scanFirstPageGsiResult = await chickenStore.scanOnePageWithGsi({ limit: 2, gsiId: 'gsi' })
+      expect(scanFirstPageGsiResult.items).toEqual([ginger, babs])
+
+      const scanSecondPageGsiResult = await chickenStore.scanOnePageWithGsi({
+        limit: 2,
+        gsiId: 'gsi',
+        exclusiveStartKey: scanFirstPageGsiResult.lastEvaluatedKey
+      })
+      expect(scanSecondPageGsiResult.items).toEqual([bunty])
+    })
+  })
+
   describe('queries', () => {
     let chickenStore: SingleEntityOperations<
       Chicken,