@twin.org/entity-storage-connector-cosmosdb 0.9.1 → 0.9.2-next.10

package/dist/es/cosmosDbEntityStorageConnector.js

@@ -1,10 +1,11 @@
 // Copyright 2024 IOTA Stiftung.
 // SPDX-License-Identifier: Apache-2.0.
 import { BulkOperationType, CosmosClient, PartitionKeyKind } from "@azure/cosmos";
+import { HealthCategory, HealthStatus } from "@twin.org/api-models";
 import { ContextIdHelper, ContextIdStore } from "@twin.org/context";
-import { BaseError, Coerce, ComponentFactory, GeneralError, Guards, HealthStatus, Is, ObjectHelper, Validation } from "@twin.org/core";
+import { BaseError, Coerce, ComponentFactory, ConflictError, GeneralError, Guards, Is, Mutex, ObjectHelper, RandomHelper, Validation } from "@twin.org/core";
 import { ComparisonOperator, EntitySchemaFactory, EntitySchemaHelper, LogicalOperator, SortDirection } from "@twin.org/entity";
-import { EntityStorageHelper } from "@twin.org/entity-storage-models";
+import { ConnectionHelper, EntityStorageHelper } from "@twin.org/entity-storage-models";
 /**
  * Class for performing entity storage operations using Cosmos DB.
  */
@@ -28,6 +29,16 @@ export class CosmosDbEntityStorageConnector {
      * @internal
      */
     static _PARTITION_KEY_VALUE = "root";
+    /**
+     * Batch chunk size for bulk write operations.
+     * @internal
+     */
+    static _BATCH_CHUNK_SIZE = 1000;
+    /**
+     * Number of bulk operation chunks to dispatch concurrently in setBatch.
+     * @internal
+     */
+    static _WRITE_CONCURRENCY = 10;
     /**
      * The name for the schema.
      * @internal
@@ -48,21 +59,26 @@ export class CosmosDbEntityStorageConnector {
      * @internal
      */
     _primaryKey;
+    /**
+     * The name of the version property, if any.
+     * @internal
+     */
+    _versionKey;
     /**
      * The configuration for the connector.
      * @internal
      */
     _config;
     /**
-     * The Cosmos DB client.
+     * Milliseconds to wait for optimistic-lock mutexes before throwing.
      * @internal
      */
-    _client;
+    _mutexTimeoutMs;
     /**
-     * Container user for the data storage.
+     * Unique identifier for this connector instance, used to track references in SharedStore.
      * @internal
      */
-    _container;
+    _instanceId;
     /**
      * Create a new instance of CosmosDbEntityStorageConnector.
      * @param options The options for the connector.
@@ -79,17 +95,18 @@ export class CosmosDbEntityStorageConnector {
         this._entitySchemaName = options.entitySchema;
         this._partitionContextIds = options.partitionContextIds;
         this._primaryKey = EntitySchemaHelper.getPrimaryKey(this._entitySchema);
+        this._versionKey = EntitySchemaHelper.findVersionProperty(this._entitySchema);
         this._config = options.config;
-        this._client = new CosmosClient({
-            endpoint: this._config.endpoint,
-            key: this._config.key,
-            connectionPolicy: {
-                enableEndpointDiscovery: !this._config.disableEndpointDiscovery
-            }
-        });
-        this._container = this._client
-            .database(this._config.databaseId)
-            .container(this._config.containerId);
+        this._mutexTimeoutMs = Coerce.integer(options.config.mutexTimeoutMs);
+        this._instanceId = RandomHelper.generateUuidV7("compact");
+    }
+    /**
+     * The component needs to be stopped when the node is closed.
+     * @param nodeLoggingComponentType The node logging component type.
+     * @returns Nothing.
+     */
+    async stop(nodeLoggingComponentType) {
+        await ConnectionHelper.closeClient("cosmosDbClients", this.createClientId(), this._instanceId, this._mutexTimeoutMs, async (client) => client.dispose());
     }
     /**
      * Initialize the Cosmos DB environment.
@@ -98,6 +115,7 @@ export class CosmosDbEntityStorageConnector {
      */
     async bootstrap(nodeLoggingComponentType) {
         const nodeLogging = ComponentFactory.getIfExists(nodeLoggingComponentType);
+        const client = await this.getClient();
         // Create the database if it does not exist
         try {
             const databaseExists = await this.databaseExists();
@@ -122,7 +140,7 @@ export class CosmosDbEntityStorageConnector {
                         databaseId: this._config.databaseId
                     }
                 });
-                await this._client.databases.create({
+                await client.databases.create({
                     id: this._config.databaseId
                 });
                 await this.waitForDatabaseExists();
@@ -165,12 +183,13 @@ export class CosmosDbEntityStorageConnector {
                         containerId: this._config.containerId
                     }
                 });
-                await this._client.database(this._config.databaseId).containers.create({
+                await client.database(this._config.databaseId).containers.create({
                     id: this._config.containerId,
                     partitionKey: {
                         kind: PartitionKeyKind.Hash,
                         paths: [`/${CosmosDbEntityStorageConnector._PARTITION_KEY}`]
-                    }
+                    },
+                    indexingPolicy: this.buildIndexingPolicy()
                 }, { offerThroughput: this._config.offerThroughput });
                 await this.waitForContainerExists();
             }
@@ -203,13 +222,12 @@ export class CosmosDbEntityStorageConnector {
      */
     async health() {
         try {
-            await this._client
-                .database(this._config.databaseId)
-                .container(this._config.containerId)
-                .read();
+            const container = await this.getContainer();
+            await container.read();
             return [
                 {
                     source: CosmosDbEntityStorageConnector.CLASS_NAME,
+                    category: HealthCategory.Connectivity,
                     status: HealthStatus.Ok,
                     description: "healthDescription",
                     data: { databaseId: this._config.databaseId, containerId: this._config.containerId }
@@ -220,6 +238,7 @@ export class CosmosDbEntityStorageConnector {
             return [
                 {
                     source: CosmosDbEntityStorageConnector.CLASS_NAME,
+                    category: HealthCategory.Connectivity,
                     status: HealthStatus.Error,
                     description: "healthDescription",
                     message: "connectionFailed",
@@ -248,11 +267,15 @@ export class CosmosDbEntityStorageConnector {
         const contextIds = await ContextIdStore.getContextIds();
         const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
         try {
+            const container = await this.getContainer();
             // No secondary index or conditions
             if (Is.empty(secondaryIndex) && !Is.arrayValue(conditions)) {
-                const { resource: item } = await this._container
+                const { resource: item } = await container
                     .item(id, partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE)
                     .read();
+                if (Is.empty(item)) {
+                    return undefined;
+                }
                 return this.itemToEntity(item);
             }
             const conditionValues = [];
@@ -286,7 +309,7 @@ export class CosmosDbEntityStorageConnector {
                 query: `SELECT * FROM c WHERE ${whereQuery.join(" AND ")}`,
                 parameters: conditionValues
             };
-            const { resources: items } = await this._container.items.query(query).fetchAll();
+            const { resources: items } = await container.items.query(query).fetchAll();
             if (items.length === 1) {
                 return this.itemToEntity(items[0]);
             }
@@ -314,25 +337,66 @@ export class CosmosDbEntityStorageConnector {
         EntityStorageHelper.validateConditions(this._entitySchema, conditions);
         const contextIds = await ContextIdStore.getContextIds();
         const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
+        const submittedVersion = Is.stringValue(this._versionKey)
+            ? ObjectHelper.propertyGet(entity, this._versionKey)
+            : undefined;
+        const hasVersionCheck = !Is.empty(this._versionKey) && !Is.empty(submittedVersion) && submittedVersion > 0;
         const prepared = EntityStorageHelper.prepareEntity(entity, this._entitySchema, undefined, {
             nullBehavior: "omit"
         });
         const id = prepared[this._primaryKey.property];
+        const optimisticMutexKey = Is.stringValue(this._versionKey)
+            ? this.buildOptimisticMutexKey(partitionKey, id)
+            : undefined;
+        if (Is.stringValue(optimisticMutexKey)) {
+            await Mutex.lock(optimisticMutexKey, {
+                throwOnTimeout: true,
+                timeoutMs: this._mutexTimeoutMs
+            });
+        }
         try {
-            if (Is.arrayValue(conditions)) {
-                const item = this._container.item(id, partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE);
+            const container = await this.getContainer();
+            let itemEtag;
+            if (Is.stringValue(this._versionKey) || Is.arrayValue(conditions)) {
+                const item = container.item(id, partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE);
                 const { resource: itemData } = await item.read();
-                if (Is.notEmpty(itemData) && !this.verifyConditions(conditions, itemData)) {
-                    return;
+                if (!Is.empty(itemData)) {
+                    if (hasVersionCheck) {
+                        const storedVersion = ObjectHelper.propertyGet(itemData, this._versionKey) ?? 0;
+                        if (storedVersion !== submittedVersion) {
+                            throw new ConflictError(CosmosDbEntityStorageConnector.CLASS_NAME, "optimisticLockFailed", id);
+                        }
+                    }
+                    if (Is.arrayValue(conditions) && !this.verifyConditions(conditions, itemData)) {
+                        if (Is.stringValue(this._versionKey)) {
+                            throw new ConflictError(CosmosDbEntityStorageConnector.CLASS_NAME, "conditionFailed", id);
+                        }
+                        return;
+                    }
+                    itemEtag = ObjectHelper.propertyGet(itemData, "_etag");
+                }
+                if (Is.stringValue(this._versionKey)) {
+                    const storedVersion = !Is.empty(itemData)
+                        ? (ObjectHelper.propertyGet(itemData, this._versionKey) ?? 0)
+                        : 0;
+                    ObjectHelper.propertySet(prepared, this._versionKey, storedVersion + 1);
                 }
             }
-            await this._container.items.upsert({
+            await container.items.upsert({
                 id,
                 [CosmosDbEntityStorageConnector._PARTITION_KEY]: partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE,
                 ...prepared
-            });
+            }, Is.stringValue(itemEtag)
+                ? { accessCondition: { type: "IfMatch", condition: itemEtag } }
+                : undefined);
         }
         catch (err) {
+            if (BaseError.isErrorName(err, ConflictError.CLASS_NAME)) {
+                throw err;
+            }
+            if (Is.object(err) && err.code === 412) {
+                throw new ConflictError(CosmosDbEntityStorageConnector.CLASS_NAME, "optimisticLockFailed", id);
+            }
             if (BaseError.isAggregateError(err)) {
                 const errors = BaseError.fromAggregate(err);
                 if (BaseError.someErrorCode(errors, "ResourceNotFoundException")) {
@@ -345,6 +409,11 @@ export class CosmosDbEntityStorageConnector {
                 id
             }, err);
         }
+        finally {
+            if (Is.stringValue(optimisticMutexKey)) {
+                Mutex.unlock(optimisticMutexKey);
+            }
+        }
     }
     /**
      * Set multiple entities in a batch.
@@ -359,15 +428,30 @@ export class CosmosDbEntityStorageConnector {
             nullBehavior: "omit"
         }));
         try {
-            await this._container.items.executeBulkOperations(preparedEntities.map(prepared => ({
-                operationType: BulkOperationType.Upsert,
-                partitionKey: partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE,
-                resourceBody: {
-                    id: prepared[this._primaryKey.property],
-                    [CosmosDbEntityStorageConnector._PARTITION_KEY]: partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE,
-                    ...prepared
+            const container = await this.getContainer();
+            const pk = partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE;
+            const chunkSize = CosmosDbEntityStorageConnector._BATCH_CHUNK_SIZE;
+            const concurrency = CosmosDbEntityStorageConnector._WRITE_CONCURRENCY;
+            const windowSize = chunkSize * concurrency;
+            for (let offset = 0; offset < preparedEntities.length; offset += windowSize) {
+                const window = preparedEntities.slice(offset, offset + windowSize);
+                const sends = [];
+                for (let j = 0; j < window.length; j += chunkSize) {
+                    const chunk = window.slice(j, j + chunkSize);
+                    sends.push((async () => {
+                        await container.items.executeBulkOperations(chunk.map(prepared => ({
+                            operationType: BulkOperationType.Upsert,
+                            partitionKey: pk,
+                            resourceBody: {
+                                id: prepared[this._primaryKey.property],
+                                [CosmosDbEntityStorageConnector._PARTITION_KEY]: pk,
+                                ...prepared
+                            }
+                        })));
+                    })());
                 }
-            })));
+                await Promise.all(sends);
+            }
         }
         catch (err) {
             throw new GeneralError(CosmosDbEntityStorageConnector.CLASS_NAME, "setBatchFailed", undefined, err);
@@ -382,10 +466,11 @@ export class CosmosDbEntityStorageConnector {
         const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
         const pk = partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE;
         try {
+            const container = await this.getContainer();
             let continuationToken;
             do {
                 const feedOptions = { maxItemCount: 100, continuationToken };
-                const { resources, continuationToken: nextToken } = await this._container.items
+                const { resources, continuationToken: nextToken } = await container.items
                     .query({
                     query: `SELECT c.id FROM c WHERE c.${CosmosDbEntityStorageConnector._PARTITION_KEY} = @pk`,
                     parameters: [{ name: "@pk", value: pk }]
@@ -398,7 +483,7 @@ export class CosmosDbEntityStorageConnector {
                         id: r.id,
                         partitionKey: pk
                     }));
-                    await this._container.items.executeBulkOperations(operations);
+                    await container.items.executeBulkOperations(operations);
                 }
             } while (Is.stringValue(continuationToken));
         }
@@ -417,17 +502,33 @@ export class CosmosDbEntityStorageConnector {
         EntityStorageHelper.validateConditions(this._entitySchema, conditions);
         const contextIds = await ContextIdStore.getContextIds();
         const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
+        const optimisticMutexKey = Is.stringValue(this._versionKey)
+            ? this.buildOptimisticMutexKey(partitionKey, id)
+            : undefined;
+        if (Is.stringValue(optimisticMutexKey)) {
+            await Mutex.lock(optimisticMutexKey, {
+                throwOnTimeout: true,
+                timeoutMs: this._mutexTimeoutMs
+            });
+        }
         try {
-            const item = this._container.item(id, partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE);
+            const container = await this.getContainer();
+            const item = container.item(id, partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE);
             const { resource: itemData } = await item.read();
             if (Is.notEmpty(itemData)) {
                 if (Is.arrayValue(conditions) && !this.verifyConditions(conditions, itemData)) {
+                    if (Is.stringValue(this._versionKey)) {
+                        throw new ConflictError(CosmosDbEntityStorageConnector.CLASS_NAME, "conditionFailed", id);
+                    }
                     return;
                 }
                 await item.delete();
             }
         }
         catch (err) {
+            if (BaseError.isErrorName(err, ConflictError.CLASS_NAME)) {
+                throw err;
+            }
             if (BaseError.fromError(err) &&
                 Is.object(err) &&
                 err.body?.code === "NotFound") {
@@ -437,6 +538,11 @@ export class CosmosDbEntityStorageConnector {
                 id
             }, err);
         }
+        finally {
+            if (Is.stringValue(optimisticMutexKey)) {
+                Mutex.unlock(optimisticMutexKey);
+            }
+        }
     }
     /**
      * Remove multiple entities by id.
@@ -448,12 +554,13 @@ export class CosmosDbEntityStorageConnector {
         const contextIds = await ContextIdStore.getContextIds();
         const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
         try {
+            const container = await this.getContainer();
             const operations = ids.map(id => ({
                 operationType: BulkOperationType.Delete,
                 id,
                 partitionKey: partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE
             }));
-            await this._container.items.executeBulkOperations(operations);
+            await container.items.executeBulkOperations(operations);
         }
         catch (err) {
             throw new GeneralError(CosmosDbEntityStorageConnector.CLASS_NAME, "removeBatchFailed", undefined, err);
@@ -474,8 +581,9 @@ export class CosmosDbEntityStorageConnector {
             data: { containerId: this._config.containerId }
         });
         try {
+            const container = await this.getContainer();
             if (await this.containerExists()) {
-                await this._container.delete();
+                await container.delete();
                 await this.waitForContainerNotExists();
             }
             await nodeLogging?.log({
@@ -515,22 +623,35 @@ export class CosmosDbEntityStorageConnector {
         EntityStorageHelper.validateSortProperties(this._entitySchema, sortProperties);
         EntityStorageHelper.validateProperties(this._entitySchema, properties);
         EntityStorageHelper.validateConditionProperties(this._entitySchema, conditions);
+        // Only the sort shapes covered by the composite indexes from buildIndexingPolicy are accepted.
+        if (Is.arrayValue(sortProperties)) {
+            const nonPrimarySorts = sortProperties.filter(sortProperty => sortProperty.property !== this._primaryKey.property);
+            if (nonPrimarySorts.length > 1) {
+                throw new GeneralError(CosmosDbEntityStorageConnector.CLASS_NAME, "sortUnsupported", {
+                    properties: sortProperties.map(sortProperty => sortProperty.property),
+                    primaryKey: this._primaryKey.property
+                });
+            }
+        }
         if (!Is.empty(limit)) {
             const validationFailures = [];
             Validation.integer("limit", limit, validationFailures, undefined, { minValue: 1 });
             Validation.asValidationError(CosmosDbEntityStorageConnector.CLASS_NAME, "query", validationFailures);
         }
         try {
+            const container = await this.getContainer();
             const returnSize = limit ?? CosmosDbEntityStorageConnector._DEFAULT_LIMIT;
             let orderByClause = "";
-            if (Array.isArray(sortProperties)) {
-                if (sortProperties.length > 1) {
-                    throw new GeneralError(CosmosDbEntityStorageConnector.CLASS_NAME, "sortSingle");
-                }
-                for (const sortProperty of sortProperties) {
+            if (Is.arrayValue(sortProperties)) {
+                // The primary key is unique so no property after its first occurrence can affect
+                // the order, truncating keeps the ORDER BY within the provisioned composite indexes.
+                const primaryKeyIndex = sortProperties.findIndex(sortProperty => sortProperty.property === this._primaryKey.property);
+                const effectiveSorts = primaryKeyIndex === -1 ? sortProperties : sortProperties.slice(0, primaryKeyIndex + 1);
+                const orderClauses = effectiveSorts.map(sortProperty => {
                     const direction = sortProperty.sortDirection === SortDirection.Ascending ? "asc" : "desc";
-                    orderByClause = `ORDER BY c.${String(sortProperty.property)} ${direction}`;
-                }
+                    return `c.${String(sortProperty.property)} ${direction}`;
+                });
+                orderByClause = `ORDER BY ${orderClauses.join(", ")}`;
             }
             const attributeNames = {};
             const attributeValues = {};
@@ -548,13 +669,13 @@ export class CosmosDbEntityStorageConnector {
                 ...Object.keys(attributeValues).map(key => ({ name: `@${key}`, value: attributeValues[key] }))
             ];
             const queryPartitionKey = partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE;
-            // For sorted queries use OFFSET LIMIT — CosmosDB continuation tokens are
+            // For sorted queries use OFFSET LIMIT - CosmosDB continuation tokens are
             // not reliably returned for ORDER BY queries across all service versions.
             // The cursor is a numeric offset encoded as a string (same as SQL connectors).
             const startIndex = Coerce.number(cursor) ?? 0;
             sql = `${baseQuery} ${orderByClause} OFFSET ${startIndex} LIMIT ${returnSize + 1}`;
             const sortedQuerySpecs = { query: sql, parameters: queryParameters };
-            const sortedResponse = await this._container.items
+            const sortedResponse = await container.items
                 .query(sortedQuerySpecs, {
                 partitionKey: queryPartitionKey,
                 maxItemCount: returnSize + 1
@@ -579,6 +700,7 @@ export class CosmosDbEntityStorageConnector {
     async count(conditions) {
         EntityStorageHelper.validateConditionProperties(this._entitySchema, conditions);
         try {
+            const container = await this.getContainer();
             const contextIds = await ContextIdStore.getContextIds();
             const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
             const attributeNames = {};
@@ -599,7 +721,7 @@ export class CosmosDbEntityStorageConnector {
                     ...Object.keys(attributeValues).map(key => ({ name: `@${key}`, value: attributeValues[key] }))
                 ]
             };
-            const { resources } = await this._container.items.query(querySpec).fetchAll();
+            const { resources } = await container.items.query(querySpec).fetchAll();
             return resources[0] ?? 0;
         }
         catch (err) {
@@ -613,10 +735,11 @@ export class CosmosDbEntityStorageConnector {
     async getPartitionContextIds() {
         const partitionContextIds = this._partitionContextIds;
         if (!Is.arrayValue(partitionContextIds)) {
-            return [];
+            return undefined;
         }
         try {
-            const { resources: partitionIds } = await this._container.items
+            const container = await this.getContainer();
+            const { resources: partitionIds } = await container.items
                 .query({
                 query: `SELECT DISTINCT VALUE c.${CosmosDbEntityStorageConnector._PARTITION_KEY} FROM c`
             })
@@ -629,6 +752,13 @@ export class CosmosDbEntityStorageConnector {
             throw new GeneralError(CosmosDbEntityStorageConnector.CLASS_NAME, "getPartitionContextIdsFailed", undefined, err);
         }
     }
+    /**
+     * Get the connector implementation version.
+     * @returns The connector implementation version.
+     */
+    connectorVersion() {
+        return 0;
+    }
     /**
      * Create the target connector for performing the migration using a temporary container.
      * @param newEntitySchema The name of the new entity schema to create the connector for.
@@ -689,22 +819,24 @@ export class CosmosDbEntityStorageConnector {
      * @internal
      */
     async bulkCopy(sourceConnector, destConnector, partitions, batchSize) {
-        let partitionList;
-        if (Is.arrayValue(partitions)) {
-            partitionList = partitions;
-        }
-        else if (Is.arrayValue(sourceConnector._partitionContextIds)) {
-            partitionList = [];
-        }
-        else {
-            partitionList = [{}];
+        // undefined → not partitioned: one pass with no partition key.
+        // []        → partitioned but empty: nothing to copy, return early.
+        // [{…}, …]  → partitioned with data: iterate over each partition.
+        if (partitions?.length === 0) {
+            return;
         }
+        const partitionList = partitions ?? [{}];
         for (let i = 0; i < partitionList.length; i++) {
-            const partitionKey = ContextIdHelper.combinedContextKey(partitionList[i], sourceConnector._partitionContextIds) ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE;
+            // Values from getPartitionContextIds are already short-form, so we join them
+            // directly rather than using combinedContextKey, which expects long-form input
+            // and calls guardAll (throwing if a registered handler rejects short-form values).
+            const partitionKey = Is.arrayValue(sourceConnector._partitionContextIds)
+                ? sourceConnector._partitionContextIds.map(k => partitionList[i][k]).join("/")
+                : CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE;
             let continuationToken;
             do {
                 const feedOptions = { maxItemCount: batchSize, continuationToken };
-                const { resources, continuationToken: nextToken } = await sourceConnector._container.items
+                const { resources, continuationToken: nextToken } = await (await sourceConnector.getContainer()).items
                     .query({
                     query: `SELECT * FROM c WHERE c.${CosmosDbEntityStorageConnector._PARTITION_KEY} = @partitionId`,
                     parameters: [{ name: "@partitionId", value: partitionKey }]
@@ -717,7 +849,7 @@ export class CosmosDbEntityStorageConnector {
                         partitionKey,
                         resourceBody: rest
                     }));
-                    await destConnector._container.items.executeBulkOperations(operations);
+                    await (await destConnector.getContainer()).items.executeBulkOperations(operations);
                 }
             } while (Is.stringValue(continuationToken));
         }
@@ -786,7 +918,7 @@ export class CosmosDbEntityStorageConnector {
         else if (Is.array(comparator.value)) {
             const dbValues = comparator.value.map(v => this.propertyToDbValue(v, type));
             if (dbValues.length === 0 && comparator.comparison === ComparisonOperator.In) {
-                // CosmosDB rejects `IN ()` — return always-false sentinel (#141).
+                // CosmosDB rejects `IN ()` - return always-false sentinel (#141).
                 return "1=0";
             }
             const arrAttributeNames = [];
@@ -945,6 +1077,16 @@ export class CosmosDbEntityStorageConnector {
     verifyConditions(conditions, obj) {
         return conditions.every(condition => ObjectHelper.propertyGet(obj, condition.property) === condition.value);
     }
+    /**
+     * Build a mutex key for optimistic-locking critical sections.
+     * @param partitionKey The resolved partition key.
+     * @param id The entity id.
+     * @returns The mutex key.
+     * @internal
+     */
+    buildOptimisticMutexKey(partitionKey, id) {
+        return `${CosmosDbEntityStorageConnector.CLASS_NAME}:optimistic:${this._config.databaseId}:${this._config.containerId}:${partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE}:${id}`;
+    }
     /**
      * Convert an entity to an item.
      * @param item The item to convert.
@@ -961,6 +1103,38 @@ export class CosmosDbEntityStorageConnector {
             "_ts"
         ]);
     }
+    /**
+     * Retrieve (or lazily create) the shared Cosmos DB client for this endpoint.
+     * @returns The shared client.
+     * @internal
+     */
+    async getClient() {
+        return ConnectionHelper.openClient("cosmosDbClients", this.createClientId(), this._instanceId, this._mutexTimeoutMs, async () => new CosmosClient({
+            endpoint: this._config.endpoint,
+            key: this._config.key,
+            connectionPolicy: {
+                enableEndpointDiscovery: !this._config.disableEndpointDiscovery
+            }
+        }));
+    }
+    /**
+     * Get the container for this connector's configured database and container.
+     * @returns The container reference.
+     * @internal
+     */
+    async getContainer() {
+        return (await this.getClient())
+            .database(this._config.databaseId)
+            .container(this._config.containerId);
+    }
+    /**
+     * Build a stable cache key for the shared client based on connection parameters.
+     * @returns The client cache key.
+     * @internal
+     */
+    createClientId() {
+        return `${this._config.endpoint}|${this._config.databaseId}|${this._config.containerId}`;
+    }
     /**
      * Check if the database exists.
      * @returns True if the database exists, false otherwise.
@@ -968,7 +1142,8 @@ export class CosmosDbEntityStorageConnector {
      */
     async databaseExists() {
         try {
-            const { resources: databaseList } = await this._client.databases.readAll().fetchAll();
+            const client = await this.getClient();
+            const { resources: databaseList } = await client.databases.readAll().fetchAll();
             return databaseList.some((db) => db.id === this._config.databaseId);
         }
         catch {
@@ -989,6 +1164,41 @@ export class CosmosDbEntityStorageConnector {
             await new Promise(resolve => setTimeout(resolve, 250));
         }
     }
+    /**
+     * Build the composite indexes needed to serve multi-property ORDER BY queries.
+     * Pairing each sortable property with the primary key in both directions covers all
+     * four direction combinations, as Cosmos DB also serves each index reversed.
+     * @returns The indexing policy for the container, or undefined if the schema has no
+     * sortable properties.
+     * @internal
+     */
+    buildIndexingPolicy() {
+        const primaryKeyPath = `/${this._primaryKey.property}`;
+        const compositeIndexes = [];
+        if (Is.arrayValue(this._entitySchema.properties)) {
+            for (const prop of this._entitySchema.properties) {
+                if (!prop.isPrimary && (Is.stringValue(prop.sortDirection) || prop.isSecondary)) {
+                    const propertyPath = `/${prop.property}`;
+                    compositeIndexes.push([
+                        { path: propertyPath, order: "ascending" },
+                        { path: primaryKeyPath, order: "ascending" }
+                    ]);
+                    compositeIndexes.push([
+                        { path: propertyPath, order: "ascending" },
+                        { path: primaryKeyPath, order: "descending" }
+                    ]);
+                }
+            }
+        }
+        return compositeIndexes.length > 0
+            ? {
+                indexingMode: "consistent",
+                automatic: true,
+                includedPaths: [{ path: "/*" }],
+                compositeIndexes
+            }
+            : undefined;
+    }
     /**
      * Check if the container exists.
      * @returns True if the container exists, false otherwise.
@@ -996,7 +1206,8 @@ export class CosmosDbEntityStorageConnector {
      */
     async containerExists() {
         try {
-            const { resources: containers } = await this._client
+            const client = await this.getClient();
+            const { resources: containers } = await client
                 .database(this._config.databaseId)
                 .containers.readAll()
                 .fetchAll();