@twin.org/entity-storage-connector-dynamodb 0.0.2-next.9 → 0.0.3-next.10

package/dist/{esm/index.mjs → es/dynamoDbEntityStorageConnector.js}

@@ -1,39 +1,44 @@
-import { waitUntilTableExists, DynamoDB, QueryCommand } from '@aws-sdk/client-dynamodb';
-import { GetCommand, PutCommand, DeleteCommand, DynamoDBDocumentClient } from '@aws-sdk/lib-dynamodb';
-import { unmarshall } from '@aws-sdk/util-dynamodb';
-import { Guards, Is, ComponentFactory, BaseError, GeneralError, Coerce, ObjectHelper, Converter } from '@twin.org/core';
-import { EntitySchemaFactory, EntitySchemaHelper, ComparisonOperator, LogicalOperator, SortDirection } from '@twin.org/entity';
-
 // Copyright 2024 IOTA Stiftung.
 // SPDX-License-Identifier: Apache-2.0.
+import { DynamoDB, QueryCommand, waitUntilTableExists } from "@aws-sdk/client-dynamodb";
+import { BatchWriteCommand, DeleteCommand, DynamoDBDocumentClient, GetCommand, PutCommand, ScanCommand } from "@aws-sdk/lib-dynamodb";
+import { unmarshall } from "@aws-sdk/util-dynamodb";
+import { ContextIdHelper, ContextIdStore } from "@twin.org/context";
+import { BaseError, Coerce, ComponentFactory, Converter, GeneralError, Guards, HealthStatus, Is, ObjectHelper } from "@twin.org/core";
+import { ComparisonOperator, EntitySchemaFactory, EntitySchemaHelper, LogicalOperator, SortDirection } from "@twin.org/entity";
 /**
  * Class for performing entity storage operations using Dynamo DB.
  */
-class DynamoDbEntityStorageConnector {
+export class DynamoDbEntityStorageConnector {
+    /**
+     * Runtime name for the class.
+     */
+    static CLASS_NAME = "DynamoDbEntityStorageConnector";
     /**
      * Limit the number of entities when finding.
      * @internal
      */
-    static _PAGE_SIZE = 40;
+    static _DEFAULT_LIMIT = 40;
     /**
      * Partition id field name.
      * @internal
      */
-    static _PARTITION_ID_NAME = "partitionId";
+    static _PARTITION_KEY = "partitionId";
     /**
      * Partition id field value.
      * @internal
      */
-    static _PARTITION_ID_VALUE = "1";
-    /**
-     * Runtime name for the class.
-     */
-    CLASS_NAME = "DynamoDbEntityStorageConnector";
+    static _PARTITION_KEY_VALUE = "root";
     /**
      * The schema for the entity.
      * @internal
      */
     _entitySchema;
+    /**
+     * The keys to use from the context ids to create partitions.
+     * @internal
+     */
+    _partitionContextIds;
     /**
      * The primary key.
      * @internal
@@ -49,23 +54,65 @@ class DynamoDbEntityStorageConnector {
      * @param options The options for the connector.
      */
     constructor(options) {
-        Guards.object(this.CLASS_NAME, "options", options);
-        Guards.stringValue(this.CLASS_NAME, "options.entitySchema", options.entitySchema);
-        Guards.object(this.CLASS_NAME, "options.config", options.config);
+        Guards.object(DynamoDbEntityStorageConnector.CLASS_NAME, "options", options);
+        Guards.stringValue(DynamoDbEntityStorageConnector.CLASS_NAME, "options.entitySchema", options.entitySchema);
+        Guards.object(DynamoDbEntityStorageConnector.CLASS_NAME, "options.config", options.config);
         options.config.authMode ??= "credentials";
         if (options.config.authMode === "credentials") {
-            Guards.stringValue(this.CLASS_NAME, "options.config.accessKeyId", options.config.accessKeyId);
-            Guards.stringValue(this.CLASS_NAME, "options.config.secretAccessKey", options.config.secretAccessKey);
+            Guards.stringValue(DynamoDbEntityStorageConnector.CLASS_NAME, "options.config.accessKeyId", options.config.accessKeyId);
+            Guards.stringValue(DynamoDbEntityStorageConnector.CLASS_NAME, "options.config.secretAccessKey", options.config.secretAccessKey);
         }
-        Guards.stringValue(this.CLASS_NAME, "options.config.region", options.config.region);
-        Guards.stringValue(this.CLASS_NAME, "options.config.tableName", options.config.tableName);
+        Guards.stringValue(DynamoDbEntityStorageConnector.CLASS_NAME, "options.config.region", options.config.region);
+        Guards.stringValue(DynamoDbEntityStorageConnector.CLASS_NAME, "options.config.tableName", options.config.tableName);
         this._entitySchema = EntitySchemaFactory.get(options.entitySchema);
+        this._partitionContextIds = options.partitionContextIds;
         this._primaryKey = EntitySchemaHelper.getPrimaryKey(this._entitySchema);
         this._config = options.config;
         this._config.endpoint = Is.stringValue(this._config.endpoint)
             ? this._config.endpoint
             : undefined;
     }
+    /**
+     * Returns the class name of the component.
+     * @returns The class name of the component.
+     */
+    className() {
+        return DynamoDbEntityStorageConnector.CLASS_NAME;
+    }
+    /**
+     * Returns the health status of the component.
+     * @returns The health status of the component.
+     */
+    async health() {
+        try {
+            const dbConnection = this.createConnection();
+            await dbConnection.describeTable({ TableName: this._config.tableName });
+            return [
+                {
+                    source: DynamoDbEntityStorageConnector.CLASS_NAME,
+                    status: HealthStatus.Ok,
+                    description: "healthDescription"
+                }
+            ];
+        }
+        catch {
+            return [
+                {
+                    source: DynamoDbEntityStorageConnector.CLASS_NAME,
+                    status: HealthStatus.Error,
+                    description: "healthDescription",
+                    message: "connectionFailed"
+                }
+            ];
+        }
+    }
+    /**
+     * Get the schema for the entities.
+     * @returns The schema for the entities.
+     */
+    getSchema() {
+        return this._entitySchema;
+    }
     /**
      * Bootstrap the component by creating and initializing any resources it needs.
      * @param nodeLoggingComponentType The node logging component type.
@@ -76,7 +123,7 @@ class DynamoDbEntityStorageConnector {
         if (!(await this.tableExists(this._config.tableName))) {
             await nodeLogging?.log({
                 level: "info",
-                source: this.CLASS_NAME,
+                source: DynamoDbEntityStorageConnector.CLASS_NAME,
                 ts: Date.now(),
                 message: "tableCreating",
                 data: {
@@ -97,11 +144,11 @@ class DynamoDbEntityStorageConnector {
                 // We always add a partition key to the table as a non optional hash key
                 // is always required when querying using sort parameters
                 tableParams.AttributeDefinitions?.push({
-                    AttributeName: DynamoDbEntityStorageConnector._PARTITION_ID_NAME,
+                    AttributeName: DynamoDbEntityStorageConnector._PARTITION_KEY,
                     AttributeType: "S"
                 });
                 tableParams.KeySchema?.push({
-                    AttributeName: DynamoDbEntityStorageConnector._PARTITION_ID_NAME,
+                    AttributeName: DynamoDbEntityStorageConnector._PARTITION_KEY,
                     KeyType: "HASH"
                 });
                 const gsi = [];
@@ -128,7 +175,7 @@ class DynamoDbEntityStorageConnector {
                                 IndexName: `${prop.property}Index`,
                                 KeySchema: [
                                     {
-                                        AttributeName: DynamoDbEntityStorageConnector._PARTITION_ID_NAME,
+                                        AttributeName: DynamoDbEntityStorageConnector._PARTITION_KEY,
                                         KeyType: "HASH"
                                     },
                                     {
@@ -160,7 +207,7 @@ class DynamoDbEntityStorageConnector {
                 });
                 await nodeLogging?.log({
                     level: "info",
-                    source: this.CLASS_NAME,
+                    source: DynamoDbEntityStorageConnector.CLASS_NAME,
                     ts: Date.now(),
                     message: "tableCreated",
                     data: {
@@ -172,7 +219,7 @@ class DynamoDbEntityStorageConnector {
                 if (BaseError.isErrorCode(err, "ResourceInUseException")) {
                     await nodeLogging?.log({
                         level: "info",
-                        source: this.CLASS_NAME,
+                        source: DynamoDbEntityStorageConnector.CLASS_NAME,
                         ts: Date.now(),
                         message: "tableExists",
                         data: {
@@ -183,7 +230,7 @@ class DynamoDbEntityStorageConnector {
                 else {
                     await nodeLogging?.log({
                         level: "error",
-                        source: this.CLASS_NAME,
+                        source: DynamoDbEntityStorageConnector.CLASS_NAME,
                         ts: Date.now(),
                         message: "tableCreateFailed",
                         error: BaseError.fromError(err),
@@ -198,7 +245,7 @@ class DynamoDbEntityStorageConnector {
         else {
             await nodeLogging?.log({
                 level: "info",
-                source: this.CLASS_NAME,
+                source: DynamoDbEntityStorageConnector.CLASS_NAME,
                 ts: Date.now(),
                 message: "tableExists",
                 data: {
@@ -208,13 +255,6 @@ class DynamoDbEntityStorageConnector {
         }
         return true;
     }
-    /**
-     * Get the schema for the entities.
-     * @returns The schema for the entities.
-     */
-    getSchema() {
-        return this._entitySchema;
-    }
     /**
      * Get an entity.
      * @param id The id of the entity to get, or the index value if secondaryIndex is set.
@@ -223,19 +263,21 @@ class DynamoDbEntityStorageConnector {
      * @returns The object if it can be found or undefined.
      */
     async get(id, secondaryIndex, conditions) {
-        Guards.stringValue(this.CLASS_NAME, "id", id);
+        Guards.stringValue(DynamoDbEntityStorageConnector.CLASS_NAME, "id", id);
+        const contextIds = await ContextIdStore.getContextIds();
+        const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
         try {
             const docClient = this.createDocClient();
             if (Is.empty(secondaryIndex) && Is.empty(conditions)) {
                 const getCommand = new GetCommand({
                     TableName: this._config.tableName,
                     Key: {
-                        [DynamoDbEntityStorageConnector._PARTITION_ID_NAME]: DynamoDbEntityStorageConnector._PARTITION_ID_VALUE,
+                        [DynamoDbEntityStorageConnector._PARTITION_KEY]: partitionKey ?? DynamoDbEntityStorageConnector._PARTITION_KEY_VALUE,
                         [this._primaryKey.property]: id
                     }
                 });
                 const response = await docClient.send(getCommand);
-                delete response.Item?.[DynamoDbEntityStorageConnector._PARTITION_ID_NAME];
+                delete response.Item?.[DynamoDbEntityStorageConnector._PARTITION_KEY];
                 return response.Item;
             }
             const finalConditions = {
@@ -257,16 +299,16 @@ class DynamoDbEntityStorageConnector {
                     });
                 }
             }
-            const queryResult = await this.internalQuery(finalConditions, undefined, undefined, undefined, 1, secondaryIndex);
+            const queryResult = await this.internalQuery(finalConditions, undefined, undefined, undefined, 1, secondaryIndex, partitionKey);
             return queryResult.entities[0];
         }
         catch (err) {
             if (BaseError.isErrorCode(err, "ResourceNotFoundException")) {
-                throw new GeneralError(this.CLASS_NAME, "tableDoesNotExist", {
-                    table: this._config.tableName
+                throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "tableDoesNotExist", {
+                    tableName: this._config.tableName
                 }, err);
             }
-            throw new GeneralError(this.CLASS_NAME, "getFailed", {
+            throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "getFailed", {
                 id
             }, err);
         }
@@ -278,7 +320,9 @@ class DynamoDbEntityStorageConnector {
      * @returns The id of the entity.
      */
     async set(entity, conditions) {
-        Guards.object(this.CLASS_NAME, "entity", entity);
+        Guards.object(DynamoDbEntityStorageConnector.CLASS_NAME, "entity", entity);
+        const contextIds = await ContextIdStore.getContextIds();
+        const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
         EntitySchemaHelper.validateEntity(entity, this.getSchema());
         const id = entity[this._primaryKey.property];
         try {
@@ -287,7 +331,7 @@ class DynamoDbEntityStorageConnector {
             const putCommand = new PutCommand({
                 TableName: this._config.tableName,
                 Item: {
-                    [DynamoDbEntityStorageConnector._PARTITION_ID_NAME]: DynamoDbEntityStorageConnector._PARTITION_ID_VALUE,
+                    [DynamoDbEntityStorageConnector._PARTITION_KEY]: partitionKey ?? DynamoDbEntityStorageConnector._PARTITION_KEY_VALUE,
                     ...entity
                 },
                 // Only set the condition expression if we have conditions to match
@@ -305,15 +349,100 @@ class DynamoDbEntityStorageConnector {
                 return;
             }
             if (BaseError.isErrorCode(err, "ResourceNotFoundException")) {
-                throw new GeneralError(this.CLASS_NAME, "tableDoesNotExist", {
+                throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "tableDoesNotExist", {
                     tableName: this._config.tableName
                 }, err);
             }
-            throw new GeneralError(this.CLASS_NAME, "setFailed", {
+            throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "setFailed", {
                 id
             }, err);
         }
     }
+    /**
+     * Set multiple entities in a batch.
+     * @param entities The entities to set.
+     * @returns Nothing.
+     */
+    async setBatch(entities) {
+        Guards.arrayValue(DynamoDbEntityStorageConnector.CLASS_NAME, "entities", entities);
+        const contextIds = await ContextIdStore.getContextIds();
+        const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
+        for (const entity of entities) {
+            EntitySchemaHelper.validateEntity(entity, this.getSchema());
+        }
+        try {
+            const docClient = this.createDocClient();
+            const chunkSize = 25;
+            for (let i = 0; i < entities.length; i += chunkSize) {
+                const chunk = entities.slice(i, i + chunkSize);
+                await docClient.send(new BatchWriteCommand({
+                    RequestItems: {
+                        [this._config.tableName]: chunk.map(entity => ({
+                            PutRequest: {
+                                Item: {
+                                    [DynamoDbEntityStorageConnector._PARTITION_KEY]: partitionKey ?? DynamoDbEntityStorageConnector._PARTITION_KEY_VALUE,
+                                    ...entity
+                                }
+                            }
+                        }))
+                    }
+                }));
+            }
+        }
+        catch (err) {
+            if (BaseError.isErrorCode(err, "ResourceNotFoundException")) {
+                throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "tableDoesNotExist", { tableName: this._config.tableName }, err);
+            }
+            throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "setBatchFailed", undefined, err);
+        }
+    }
+    /**
+     * Empty the entity storage.
+     * @returns Nothing.
+     */
+    async empty() {
+        try {
+            const contextIds = await ContextIdStore.getContextIds();
+            const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
+            const pKey = partitionKey ?? DynamoDbEntityStorageConnector._PARTITION_KEY_VALUE;
+            const docClient = this.createDocClient();
+            const chunkSize = 25;
+            let exclusiveStartKey;
+            do {
+                const scanResult = await docClient.send(new ScanCommand({
+                    TableName: this._config.tableName,
+                    FilterExpression: "#partitionId = :partitionId",
+                    ExpressionAttributeNames: {
+                        "#partitionId": DynamoDbEntityStorageConnector._PARTITION_KEY
+                    },
+                    ExpressionAttributeValues: {
+                        ":partitionId": pKey
+                    },
+                    ExclusiveStartKey: exclusiveStartKey
+                }));
+                const items = scanResult.Items ?? [];
+                for (let i = 0; i < items.length; i += chunkSize) {
+                    const chunk = items.slice(i, i + chunkSize);
+                    await docClient.send(new BatchWriteCommand({
+                        RequestItems: {
+                            [this._config.tableName]: chunk.map((item) => ({
+                                DeleteRequest: {
+                                    Key: {
+                                        [DynamoDbEntityStorageConnector._PARTITION_KEY]: item[DynamoDbEntityStorageConnector._PARTITION_KEY],
+                                        [this._primaryKey.property]: item[this._primaryKey.property]
+                                    }
+                                }
+                            }))
+                        }
+                    }));
+                }
+                exclusiveStartKey = scanResult.LastEvaluatedKey;
+            } while (exclusiveStartKey);
+        }
+        catch (err) {
+            throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "emptyFailed", undefined, err);
+        }
+    }
     /**
      * Remove the entity.
      * @param id The id of the entity to remove.
@@ -321,14 +450,16 @@ class DynamoDbEntityStorageConnector {
      * @returns Nothing.
      */
     async remove(id, conditions) {
-        Guards.stringValue(this.CLASS_NAME, "id", id);
+        Guards.stringValue(DynamoDbEntityStorageConnector.CLASS_NAME, "id", id);
+        const contextIds = await ContextIdStore.getContextIds();
+        const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
         try {
             const docClient = this.createDocClient();
             const { conditionExpression, attributeNames, attributeValues } = this.buildConditionExpression(conditions);
             const deleteCommand = new DeleteCommand({
                 TableName: this._config.tableName,
                 Key: {
-                    [DynamoDbEntityStorageConnector._PARTITION_ID_NAME]: DynamoDbEntityStorageConnector._PARTITION_ID_VALUE,
+                    [DynamoDbEntityStorageConnector._PARTITION_KEY]: partitionKey ?? DynamoDbEntityStorageConnector._PARTITION_KEY_VALUE,
                     [this._primaryKey.property]: id
                 },
                 ConditionExpression: conditionExpression,
@@ -342,38 +473,134 @@ class DynamoDbEntityStorageConnector {
                 return;
             }
             if (BaseError.isErrorCode(err, "ResourceNotFoundException")) {
-                throw new GeneralError(this.CLASS_NAME, "tableDoesNotExist", {
-                    table: this._config.tableName
+                throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "tableDoesNotExist", {
+                    tableName: this._config.tableName
                 }, err);
             }
-            throw new GeneralError(this.CLASS_NAME, "removeFailed", {
+            throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "removeFailed", {
                 id
             }, err);
         }
     }
+    /**
+     * Remove multiple entities by their IDs in a batch.
+     * @param ids The ids of the entities to remove.
+     * @returns Nothing.
+     */
+    async removeBatch(ids) {
+        Guards.arrayValue(DynamoDbEntityStorageConnector.CLASS_NAME, "ids", ids);
+        const contextIds = await ContextIdStore.getContextIds();
+        const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
+        try {
+            const docClient = this.createDocClient();
+            const chunkSize = 25;
+            const primaryKeyProperty = this._primaryKey.property;
+            for (let i = 0; i < ids.length; i += chunkSize) {
+                const chunk = ids.slice(i, i + chunkSize);
+                await docClient.send(new BatchWriteCommand({
+                    RequestItems: {
+                        [this._config.tableName]: chunk.map(id => ({
+                            DeleteRequest: {
+                                Key: {
+                                    [primaryKeyProperty]: id,
+                                    [DynamoDbEntityStorageConnector._PARTITION_KEY]: partitionKey ?? DynamoDbEntityStorageConnector._PARTITION_KEY_VALUE
+                                }
+                            }
+                        }))
+                    }
+                }));
+            }
+        }
+        catch (err) {
+            throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "removeBatchFailed", undefined, err);
+        }
+    }
+    /**
+     * Teardown the entity storage by deleting the underlying table.
+     * @param nodeLoggingComponentType The node logging component type.
+     * @returns True if the teardown process was successful.
+     */
+    async teardown(nodeLoggingComponentType) {
+        const nodeLogging = ComponentFactory.getIfExists(nodeLoggingComponentType);
+        await nodeLogging?.log({
+            level: "info",
+            source: DynamoDbEntityStorageConnector.CLASS_NAME,
+            ts: Date.now(),
+            message: "tableDeleting",
+            data: { tableName: this._config.tableName }
+        });
+        try {
+            const dbConnection = this.createConnection();
+            await dbConnection.deleteTable({ TableName: this._config.tableName });
+            await nodeLogging?.log({
+                level: "info",
+                source: DynamoDbEntityStorageConnector.CLASS_NAME,
+                ts: Date.now(),
+                message: "tableDeleted",
+                data: { tableName: this._config.tableName }
+            });
+            return true;
+        }
+        catch (err) {
+            await nodeLogging?.log({
+                level: "error",
+                source: DynamoDbEntityStorageConnector.CLASS_NAME,
+                ts: Date.now(),
+                message: "teardownFailed",
+                error: BaseError.fromError(err)
+            });
+            return false;
+        }
+    }
     /**
      * Find all the entities which match the conditions.
      * @param conditions The conditions to match for the entities.
      * @param sortProperties The optional sort order.
      * @param properties The optional properties to return, defaults to all.
-     * @param cursor The cursor to request the next page of entities.
-     * @param pageSize The suggested number of entities to return in each chunk, in some scenarios can return a different amount.
+     * @param cursor The cursor to request the next chunk of entities.
+     * @param limit The suggested number of entities to return in each chunk, in some scenarios can return a different amount.
      * @returns All the entities for the storage matching the conditions,
      * and a cursor which can be used to request more entities.
      */
-    async query(conditions, sortProperties, properties, cursor, pageSize) {
-        return this.internalQuery(conditions, sortProperties, properties, cursor, pageSize);
+    async query(conditions, sortProperties, properties, cursor, limit) {
+        const contextIds = await ContextIdStore.getContextIds();
+        const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
+        return this.internalQuery(conditions, sortProperties, properties, cursor, limit, undefined, partitionKey);
     }
     /**
-     * Delete the table.
-     * @returns Nothing.
+     * Count all the entities which match the conditions.
+     * @returns The total count of entities in the storage.
      */
-    async tableDelete() {
+    async count() {
         try {
+            const contextIds = await ContextIdStore.getContextIds();
+            const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
             const dbConnection = this.createConnection();
-            await dbConnection.deleteTable({ TableName: this._config.tableName });
+            let total = 0;
+            let exclusiveStartKey;
+            do {
+                const result = await dbConnection.send(new QueryCommand({
+                    TableName: this._config.tableName,
+                    Select: "COUNT",
+                    KeyConditionExpression: "#partitionId = :partitionId",
+                    ExpressionAttributeNames: {
+                        "#partitionId": DynamoDbEntityStorageConnector._PARTITION_KEY
+                    },
+                    ExpressionAttributeValues: {
+                        ":partitionId": {
+                            S: partitionKey ?? DynamoDbEntityStorageConnector._PARTITION_KEY_VALUE
+                        }
+                    },
+                    ExclusiveStartKey: exclusiveStartKey
+                }));
+                total += result.Count ?? 0;
+                exclusiveStartKey = result.LastEvaluatedKey;
+            } while (exclusiveStartKey);
+            return total;
+        }
+        catch (err) {
+            throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "countFailed", undefined, err);
         }
-        catch { }
     }
     /**
      * Create the parameters for a query.
@@ -492,7 +719,7 @@ class DynamoDbEntityStorageConnector {
         else if (comparator.comparison === ComparisonOperator.In) {
             return `${propName} IN ${attributeName}`;
         }
-        throw new GeneralError(this.CLASS_NAME, "comparisonNotSupported", {
+        throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "comparisonNotSupported", {
             comparison: comparator.comparison
         });
     }
@@ -529,7 +756,9 @@ class DynamoDbEntityStorageConnector {
         else if (operator === LogicalOperator.Or) {
             return "OR";
         }
-        throw new GeneralError(this.CLASS_NAME, "conditionalNotSupported", { operator });
+        throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "conditionalNotSupported", {
+            operator
+        });
     }
     /**
      * Format a value to insert into DB.
@@ -585,6 +814,9 @@ class DynamoDbEntityStorageConnector {
      * @internal
      */
     createConnectionConfig() {
+        const requestHandler = {
+            requestTimeout: this._config.connectionTimeoutMs
+        };
         if (Is.stringValue(this._config.secretAccessKey) &&
             Is.stringValue(this._config.accessKeyId) &&
             this._config.authMode === "credentials") {
@@ -594,12 +826,14 @@ class DynamoDbEntityStorageConnector {
                     secretAccessKey: this._config.secretAccessKey
                 },
                 endpoint: this._config.endpoint,
-                region: this._config.region
+                region: this._config.region,
+                requestHandler
             };
         }
         return {
             endpoint: this._config.endpoint,
-            region: this._config.region
+            region: this._config.region,
+            requestHandler
         };
     }
     /**
@@ -623,16 +857,17 @@ class DynamoDbEntityStorageConnector {
      * @param conditions The conditions to match for the entities.
      * @param sortProperties The optional sort order.
      * @param properties The optional properties to return, defaults to all.
-     * @param cursor The cursor to request the next page of entities.
-     * @param pageSize The suggested number of entities to return in each chunk, in some scenarios can return a different amount.
+     * @param cursor The cursor to request the next chunk of entities.
+     * @param limit The suggested number of entities to return in each chunk, in some scenarios can return a different amount.
      * @param secondaryIndex The secondary index to use for the query.
+     * @param partitionKey The partition key to use for the query.
      * @returns All the entities for the storage matching the conditions,
      * and a cursor which can be used to request more entities.
      * @internal
      */
-    async internalQuery(conditions, sortProperties, properties, cursor, pageSize, secondaryIndex) {
+    async internalQuery(conditions, sortProperties, properties, cursor, limit, secondaryIndex, partitionKey) {
         try {
-            const returnSize = pageSize ?? DynamoDbEntityStorageConnector._PAGE_SIZE;
+            const returnSize = limit ?? DynamoDbEntityStorageConnector._DEFAULT_LIMIT;
             let indexName = Is.stringValue(secondaryIndex)
                 ? `${secondaryIndex}Index`
                 : undefined;
@@ -641,7 +876,7 @@ class DynamoDbEntityStorageConnector {
             let scanAscending = true;
             if (Is.arrayValue(sortProperties)) {
                 if (sortProperties.length > 1) {
-                    throw new GeneralError(this.CLASS_NAME, "sortSingle");
+                    throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "sortSingle");
                 }
                 for (const sortProperty of sortProperties) {
                     const propertySchema = this._entitySchema.properties?.find(e => e.property === sortProperty.property);
@@ -649,7 +884,7 @@ class DynamoDbEntityStorageConnector {
                         (!propertySchema.isPrimary &&
                             !propertySchema.isSecondary &&
                             Is.empty(propertySchema.sortDirection))) {
-                        throw new GeneralError(this.CLASS_NAME, "sortNotIndexed", {
+                        throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "sortNotIndexed", {
                             property: sortProperty.property
                         });
                     }
@@ -661,8 +896,8 @@ class DynamoDbEntityStorageConnector {
             }
             const attributeNames = { "#partitionId": "partitionId" };
             const attributeValues = {
-                [`:${DynamoDbEntityStorageConnector._PARTITION_ID_NAME}`]: {
-                    S: DynamoDbEntityStorageConnector._PARTITION_ID_VALUE
+                [`:${DynamoDbEntityStorageConnector._PARTITION_KEY}`]: {
+                    S: partitionKey ?? DynamoDbEntityStorageConnector._PARTITION_KEY_VALUE
                 }
             };
             const expressions = this.buildQueryParameters("", conditions, attributeNames, attributeValues, secondaryIndex);
@@ -692,7 +927,7 @@ class DynamoDbEntityStorageConnector {
             if (Is.arrayValue(results.Items)) {
                 entities = results.Items.map(item => {
                     const unmarshalled = unmarshall(item);
-                    delete unmarshalled[DynamoDbEntityStorageConnector._PARTITION_ID_NAME];
+                    delete unmarshalled[DynamoDbEntityStorageConnector._PARTITION_KEY];
                     return unmarshalled;
                 });
             }
@@ -705,11 +940,11 @@ class DynamoDbEntityStorageConnector {
         }
         catch (err) {
             if (BaseError.isErrorCode(err, "ResourceNotFoundException")) {
-                throw new GeneralError(this.CLASS_NAME, "tableDoesNotExist", {
-                    table: this._config.tableName
+                throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "tableDoesNotExist", {
+                    tableName: this._config.tableName
                 }, err);
             }
-            throw new GeneralError(this.CLASS_NAME, "queryFailed", undefined, err);
+            throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "queryFailed", undefined, err);
         }
     }
     /**
@@ -728,7 +963,7 @@ class DynamoDbEntityStorageConnector {
             for (const c of conditions) {
                 const schemaProp = this._entitySchema.properties?.find(p => p.property === c.property);
                 if (Is.undefined(schemaProp)) {
-                    throw new GeneralError(this.CLASS_NAME, "propertyNotFound", {
+                    throw new GeneralError(DynamoDbEntityStorageConnector.CLASS_NAME, "propertyNotFound", {
                         property: c.property
                     });
                 }
@@ -745,5 +980,4 @@ class DynamoDbEntityStorageConnector {
         return { conditionExpression, attributeNames, attributeValues };
     }
 }
-
-export { DynamoDbEntityStorageConnector };
+//# sourceMappingURL=dynamoDbEntityStorageConnector.js.map