diff --git a/src/ddb.ts b/src/ddb.ts index 69249ff..e04b723 100644 --- a/src/ddb.ts +++ b/src/ddb.ts @@ -1,5 +1,6 @@ import {DescribeTableCommand} from '@aws-sdk/client-dynamodb'; import {BatchWriteCommand, ScanCommand} from '@aws-sdk/lib-dynamodb'; +import type {DynamoDBDocument} from '@aws-sdk/lib-dynamodb'; import type {DynamoDBClient} from '@aws-sdk/client-dynamodb'; import type {DynamoDBDocumentClient} from '@aws-sdk/lib-dynamodb'; import type { @@ -15,7 +16,10 @@ export type Credentials = { sessionToken: string; }; -export function scan(params: ScanCommandInput, client: DynamoDBClient): Promise { +export function scan( + params: ScanCommandInput, + client: DynamoDBClient | DynamoDBDocument +): Promise { const command = new ScanCommand(params); // @ts-ignore @@ -24,9 +28,10 @@ export function scan(params: ScanCommandInput, client: DynamoDBClient): Promise< export async function getTableItemsCount( tableName: string, - client: DynamoDBClient + client: DynamoDBClient | DynamoDBDocument ): Promise { const command = new DescribeTableCommand({TableName: tableName}); + // @ts-ignore const resp = await client.send(command); return resp.Table!.ItemCount!; diff --git a/src/parallel-scan-stream.ts b/src/parallel-scan-stream.ts index 907c74a..73c044b 100644 --- a/src/parallel-scan-stream.ts +++ b/src/parallel-scan-stream.ts @@ -3,7 +3,7 @@ import cloneDeep from 'lodash.clonedeep'; import times from 'lodash.times'; import chunk from 'lodash.chunk'; import getDebugger from 'debug'; -import type {ScanCommandInput} from '@aws-sdk/lib-dynamodb'; +import type {DynamoDBDocument, ScanCommandInput} from '@aws-sdk/lib-dynamodb'; import type {ScanCommandOutput} from '@aws-sdk/lib-dynamodb'; import type {DynamoDBClient} from '@aws-sdk/client-dynamodb'; import type {Credentials} from './ddb'; @@ -30,7 +30,7 @@ export async function parallelScanAsStream( chunkSize: number; highWaterMark?: number; credentials?: Credentials; - client?: DynamoDBClient; + client?: DynamoDBClient | DynamoDBDocument; } ): Promise { const ddbClient = client ?? ddbv3Client(credentials); @@ -91,7 +91,7 @@ async function getItemsFromSegment({ segmentIndex: number; chunkSize: number; blocker: Blocker; - client: DynamoDBClient; + client: DynamoDBClient | DynamoDBDocument; }): Promise { let segmentItems: ScanCommandOutput['Items'] = []; let ExclusiveStartKey: ScanCommandInput['ExclusiveStartKey']; diff --git a/src/parallel-scan.ts b/src/parallel-scan.ts index aade77b..0d8d000 100644 --- a/src/parallel-scan.ts +++ b/src/parallel-scan.ts @@ -4,6 +4,7 @@ import getDebugger from 'debug'; import type {ScanCommandInput, ScanCommandOutput} from '@aws-sdk/lib-dynamodb'; import type {Credentials} from './ddb'; import type {DynamoDBClient} from '@aws-sdk/client-dynamodb'; +import type {DynamoDBDocument} from '@aws-sdk/lib-dynamodb'; import {getTableItemsCount, scan} from './ddb'; import {ddbv3Client} from './clients'; @@ -19,7 +20,7 @@ export async function parallelScan( concurrency, credentials, client, - }: {concurrency: number; credentials?: Credentials; client?: DynamoDBClient} + }: {concurrency: number; credentials?: Credentials; client?: DynamoDBClient | DynamoDBDocument} ): Promise { const ddbClient = client ?? ddbv3Client(credentials); totalTableItemsCount = await getTableItemsCount(scanParams.TableName!, ddbClient); @@ -55,7 +56,7 @@ async function getItemsFromSegment( concurrency, segmentIndex, client, - }: {concurrency: number; segmentIndex: number; client: DynamoDBClient} + }: {concurrency: number; segmentIndex: number; client: DynamoDBClient | DynamoDBDocument} ): Promise { const segmentItems: ScanCommandOutput['Items'] = []; let ExclusiveStartKey: ScanCommandInput['ExclusiveStartKey'];