Merge pull request #254 from shelfio/feraure/INT-11-fix-types

INT-11 Align types
shelfio · Mar 25, 2024 · ca5f281 · ca5f281
2 parents 98e8782 + d63dd28
commit ca5f281
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 7 deletions.
diff --git 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<ScanCommandOutput> {
+export function scan(
+  params: ScanCommandInput,
+  client: DynamoDBClient | DynamoDBDocument
+): Promise<ScanCommandOutput> {
   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<number> {
   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
@@ -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<Readable> {
   const ddbClient = client ?? ddbv3Client(credentials);
@@ -91,7 +91,7 @@ async function getItemsFromSegment({
   segmentIndex: number;
   chunkSize: number;
   blocker: Blocker;
-  client: DynamoDBClient;
+  client: DynamoDBClient | DynamoDBDocument;
 }): Promise<void> {
   let segmentItems: ScanCommandOutput['Items'] = [];
   let ExclusiveStartKey: ScanCommandInput['ExclusiveStartKey'];

diff --git 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<ScanCommandOutput['Items']> {
   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<ScanCommandOutput['Items']> {
   const segmentItems: ScanCommandOutput['Items'] = [];
   let ExclusiveStartKey: ScanCommandInput['ExclusiveStartKey'];