feat: allow scan limit override (#866)

## Problem

Find Single Row action does a full table scan which get very expensive
for high volume pipes and large tiles.

## Solution

For some high volume cases, there is no need to scan the entire table
but just the most recent ones.

This PR allows the overriding of the scanned limit by modifying the step
config.

## Changes

- Adding `stepConfig` as an optional param to `getTableRows` function to
allow limit the number rows scanned.
- Adding `order` as a param to `getTableRows` which determines whether
to traverse the createdAt index by ascending or descending order.
- Prevent clearing the entire stepConfig on test step, we should only be
deleting the `templateConfig` key in the step config

Author: pregnantboy

package-lock.json

@@ -10,6 +10,7 @@ import {
   generateMockTableRowData,
 } from '@/graphql/__tests__/mutations/tiles/table.mock'
 import { createTableRow, TableRowFilter } from '@/models/dynamodb/table-row'
+import * as tableFunctions from '@/models/dynamodb/table-row/functions'
 import TableColumnMetadata from '@/models/table-column-metadata'
 import TableMetadata from '@/models/table-metadata'
 import User from '@/models/user'
@@ -18,6 +19,24 @@ import Context from '@/types/express/context'
 import tiles from '../..'
 import findSingleRowAction from '../../actions/find-single-row'
 
+const mocks = vi.hoisted(() => ({
+  stepQueryResult: vi.fn().mockResolvedValue({
+    config: {
+      adminOverride: {
+        tileScanLimit: 100,
+      },
+    },
+  }),
+}))
+
+vi.mock('@/models/step', () => ({
+  default: {
+    query: vi.fn().mockReturnThis(),
+    findById: vi.fn().mockReturnThis(),
+    throwIfNotFound: mocks.stepQueryResult,
+  },
+}))
+
 describe('tiles create row action', () => {
   let context: Context
   let dummyTable: TableMetadata
@@ -107,4 +126,37 @@ describe('tiles create row action', () => {
       .where({ table_id: $.step.parameters.tableId })
     await expect(findSingleRowAction.run($)).rejects.toThrow(StepError)
   })
+
+  it('should call getTableRows with scan limit if exists in step config', async () => {
+    const getTableRowsSpy = vi
+      .spyOn(tableFunctions, 'getTableRows')
+      .mockResolvedValueOnce({
+        rows: [],
+        stringifiedCursor: undefined,
+      })
+    await findSingleRowAction.run($)
+    expect(getTableRowsSpy).toHaveBeenCalledWith({
+      tableId: $.step.parameters.tableId,
+      filters: $.step.parameters.filters,
+      scanLimit: 100,
+      order: 'asc',
+    })
+  })
+
+  it('should call getTableRows with scan limit and order', async () => {
+    const getTableRowsSpy = vi
+      .spyOn(tableFunctions, 'getTableRows')
+      .mockResolvedValueOnce({
+        rows: [],
+        stringifiedCursor: undefined,
+      })
+    $.step.parameters.returnLastRow = true
+    await findSingleRowAction.run($)
+    expect(getTableRowsSpy).toHaveBeenCalledWith({
+      tableId: $.step.parameters.tableId,
+      filters: $.step.parameters.filters,
+      scanLimit: 100,
+      order: 'desc',
+    })
+  })
 })

packages/backend/src/apps/tiles/__tests__/actions/find-single-row.itest.ts

@@ -8,6 +8,7 @@ import {
   TableRowFilter,
   TableRowFilterOperator,
 } from '@/models/dynamodb/table-row'
+import Step from '@/models/step'
 import TableCollaborator from '@/models/table-collaborators'
 import TableColumnMetadata from '@/models/table-column-metadata'
 
@@ -149,6 +150,7 @@ const action: IRawAction = {
       returnLastRow: boolean | undefined
     }
 
+    const step = await Step.query().findById($.step.id).throwIfNotFound()
     /**
      * Check for columns first, there will not be any columns if the tile has been deleted.
      */
@@ -175,13 +177,16 @@ const action: IRawAction = {
         $.app.name,
       )
     }
+    // Retrieve the manual scan limit override, converting it to a number.
+    // If the conversion results in NaN, we set scanLimit to undefined.
+    const scanLimitRaw = +step.config?.adminOverride?.tileScanLimit
+    const scanLimit = isNaN(scanLimitRaw) ? undefined : scanLimitRaw
 
-    /**
-     * When columnIds are not provided, we only return rowId
-     */
     const { rows } = await getTableRows({
       tableId,
       filters,
+      order: returnLastRow ? 'desc' : 'asc',
+      scanLimit,
     })
 
     if (!rows || !rows.length) {
@@ -192,9 +197,7 @@ const action: IRawAction = {
       })
       return
     }
-    const rowIdToUse = returnLastRow
-      ? rows[rows.length - 1].rowId
-      : rows[0].rowId
+    const rowIdToUse = rows[0].rowId
 
     /**
      * We use raw row data instead of mapped column names as we want them to

packages/backend/src/apps/tiles/actions/find-single-row/index.ts

@@ -1,3 +1,5 @@
+import { raw } from 'objection'
+
 import testStep from '@/services/test-step'
 
 import type { MutationResolvers } from '../__generated__/types.generated'
@@ -29,7 +31,8 @@ const executeStep: MutationResolvers['executeStep'] = async (
   if (!executionStep.isFailed) {
     await stepToTest.$query().patch({
       status: 'completed',
-      config: {}, // clear template step config when step is tested successfully
+      // clear templateConfig in config when step is tested successfully
+      config: raw(`config - 'templateConfig'`),
     })
   }
 

packages/backend/src/graphql/mutations/execute-step.ts

@@ -90,6 +90,49 @@ describe('dynamodb table row functions', () => {
       expect(Object.keys(rows[0])).not.toContain('randomcol')
     })
 
+    describe('get table rows with scan limit', () => {
+      const TEN_THOUSAND = 10000
+      // i tried using beforeAll hook to cut down the setup time, but couldnt get it to work for this nested describe
+      beforeEach(async () => {
+        const dataArray = []
+        for (let i = 0; i < TEN_THOUSAND; i++) {
+          dataArray.push(
+            generateMockTableRowData({ columnIds: dummyColumnIds }),
+          )
+        }
+        await createTableRows({ tableId: dummyTable.id, dataArray })
+      })
+      it(
+        'should be able to paginate and get a large number of rows',
+        async () => {
+          const { rows } = await getTableRows({
+            tableId: dummyTable.id,
+            columnIds: dummyColumnIds,
+          })
+          expect(rows).toHaveLength(TEN_THOUSAND)
+        },
+        {
+          timeout: 20000,
+        },
+      )
+
+      it(
+        'should be able to paginate and get a large number of rows with a scan limit',
+        async () => {
+          const SCAN_LIMIT = 5789
+          const { rows } = await getTableRows({
+            tableId: dummyTable.id,
+            columnIds: dummyColumnIds,
+            scanLimit: SCAN_LIMIT,
+          })
+          expect(rows).toHaveLength(SCAN_LIMIT)
+        },
+        {
+          timeout: 20000,
+        },
+      )
+    })
+
     it('should get relevant rows based on a single filter', async () => {
       const dataArray = []
       for (let i = 0; i < 1000; i++) {

packages/backend/src/models/dynamodb/__tests__/table-row/functions.itest.ts

@@ -293,20 +293,31 @@ export const getTableRows = async ({
   tableId,
   columnIds,
   filters,
+  order = 'asc',
   stringifiedCursor,
+  scanLimit,
 }: {
   tableId: string
   columnIds?: string[]
   filters?: TableRowFilter[]
+  order?: 'asc' | 'desc'
   /**
    * if stringifiedCursor is 'start', we will fetch the first page of results
    * if undefined, we will auto-paginate
    */
   stringifiedCursor?: string | 'start'
+  /**
+   * Optional limit on the total number of rows scanned.
+   */
+  scanLimit?: number
 }): Promise<{
   rows: TableRowOutput[]
   stringifiedCursor?: string
 }> => {
+  if (stringifiedCursor && scanLimit) {
+    throw new Error('stringifiedCursor and scanLimit cannot both be provided')
+  }
+
   try {
     // need to use ProjectionExpression to select nested attributes
     const { ProjectionExpression, ExpressionAttributeNames } =
@@ -316,6 +327,7 @@ export const getTableRows = async ({
         indexUsed: 'byCreatedAt',
       })
     const tableRows = []
+    let remainingScanLimit = scanLimit ?? Infinity
     let cursor: any =
       stringifiedCursor && stringifiedCursor !== 'start'
         ? JSON.parse(stringifiedCursor)
@@ -326,34 +338,39 @@ export const getTableRows = async ({
         addFiltersToQuery(query, filters)
       }
       const response = await query.go({
-        order: 'asc',
+        order,
         pages: 'all', // this is ignored, we need to paginate manually
         cursor,
         params: {
           ProjectionExpression,
           ExpressionAttributeNames,
+          Limit: remainingScanLimit,
         },
         // use data:'raw' to bypass electrodb formatting, since we're using ProjectionExpression to select nested attributes
         // ref: https://electrodb.dev/en/queries/get/#execution-options
         data: 'raw',
         pager: 'raw',
         ignoreOwnership: true,
       })
+
       // need to explicitly cast to DynamoDB's raw output because of the 'raw' option
       const data = response.data as unknown as QueryCommandOutput & {
         Items: TableRowOutput[]
       }
       tableRows.push(...data.Items)
+      remainingScanLimit -= data.ScannedCount
       cursor = data.LastEvaluatedKey
       // loop only if cursor is
-    } while (cursor && !stringifiedCursor)
+    } while (cursor && !stringifiedCursor && remainingScanLimit > 0)
 
     return {
       rows: tableRows.map((row) => ({
         ...row,
         data: row.data || {}, // data can be undefined if values are empty
       })),
-      stringifiedCursor: cursor ? JSON.stringify(cursor) : undefined,
+      stringifiedCursor:
+        cursor && stringifiedCursor ? JSON.stringify(cursor) : undefined,
+      // if no cursor was passed in, we should not return
     }
   } catch (e: unknown) {
     handleDynamoDBError(e)

packages/backend/src/models/dynamodb/table-row/functions.ts

@@ -1,11 +1,10 @@
-import type { IJSONObject, IStep } from '@plumber/types'
+import type { IJSONObject, IStep, IStepConfig } from '@plumber/types'
 
 import { type StaticHookArguments, ValidationError } from 'objection'
 import { URL } from 'url'
 
 import apps from '@/apps'
 import appConfig from '@/config/app'
-import type { StepConfig } from '@/graphql/__generated__/types.generated'
 
 import Base from './base'
 import Connection from './connection'
@@ -29,7 +28,7 @@ class Step extends Base {
   connection?: Connection
   flow: Flow
   executionSteps: ExecutionStep[]
-  config: StepConfig
+  config: IStepConfig
 
   static tableName = 'steps'
 

packages/backend/src/models/step.ts

@@ -135,6 +135,7 @@ export interface IExecution {
 
 export interface IStepConfig {
   templateConfig?: IStepTemplateConfig
+  adminOverride?: IJSONObject
 }
 
 export interface IStepTemplateConfig {