@twin.org/entity-storage-connector-cosmosdb 0.9.2-next.1 → 0.9.2-next.11

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,10 +1,11 @@
1
1
  // Copyright 2024 IOTA Stiftung.
2
2
  // SPDX-License-Identifier: Apache-2.0.
3
3
  import { BulkOperationType, CosmosClient, PartitionKeyKind } from "@azure/cosmos";
4
+ import { HealthCategory, HealthStatus } from "@twin.org/api-models";
4
5
  import { ContextIdHelper, ContextIdStore } from "@twin.org/context";
5
- import { BaseError, Coerce, ComponentFactory, GeneralError, Guards, HealthStatus, Is, ObjectHelper, Validation } from "@twin.org/core";
6
+ import { BaseError, Coerce, ComponentFactory, ConflictError, GeneralError, Guards, Is, Mutex, ObjectHelper, RandomHelper, Validation } from "@twin.org/core";
6
7
  import { ComparisonOperator, EntitySchemaFactory, EntitySchemaHelper, LogicalOperator, SortDirection } from "@twin.org/entity";
7
- import { EntityStorageHelper } from "@twin.org/entity-storage-models";
8
+ import { ConnectionHelper, EntityStorageHelper } from "@twin.org/entity-storage-models";
8
9
  /**
9
10
  * Class for performing entity storage operations using Cosmos DB.
10
11
  */
@@ -28,6 +29,16 @@ export class CosmosDbEntityStorageConnector {
28
29
  * @internal
29
30
  */
30
31
  static _PARTITION_KEY_VALUE = "root";
32
+ /**
33
+ * Batch chunk size for bulk write operations.
34
+ * @internal
35
+ */
36
+ static _BATCH_CHUNK_SIZE = 1000;
37
+ /**
38
+ * Number of bulk operation chunks to dispatch concurrently in setBatch.
39
+ * @internal
40
+ */
41
+ static _WRITE_CONCURRENCY = 10;
31
42
  /**
32
43
  * The name for the schema.
33
44
  * @internal
@@ -48,21 +59,26 @@ export class CosmosDbEntityStorageConnector {
48
59
  * @internal
49
60
  */
50
61
  _primaryKey;
62
+ /**
63
+ * The name of the version property, if any.
64
+ * @internal
65
+ */
66
+ _versionKey;
51
67
  /**
52
68
  * The configuration for the connector.
53
69
  * @internal
54
70
  */
55
71
  _config;
56
72
  /**
57
- * The Cosmos DB client.
73
+ * Milliseconds to wait for optimistic-lock mutexes before throwing.
58
74
  * @internal
59
75
  */
60
- _client;
76
+ _mutexTimeoutMs;
61
77
  /**
62
- * Container user for the data storage.
78
+ * Unique identifier for this connector instance, used to track references in SharedStore.
63
79
  * @internal
64
80
  */
65
- _container;
81
+ _instanceId;
66
82
  /**
67
83
  * Create a new instance of CosmosDbEntityStorageConnector.
68
84
  * @param options The options for the connector.
@@ -79,17 +95,18 @@ export class CosmosDbEntityStorageConnector {
79
95
  this._entitySchemaName = options.entitySchema;
80
96
  this._partitionContextIds = options.partitionContextIds;
81
97
  this._primaryKey = EntitySchemaHelper.getPrimaryKey(this._entitySchema);
98
+ this._versionKey = EntitySchemaHelper.findVersionProperty(this._entitySchema);
82
99
  this._config = options.config;
83
- this._client = new CosmosClient({
84
- endpoint: this._config.endpoint,
85
- key: this._config.key,
86
- connectionPolicy: {
87
- enableEndpointDiscovery: !this._config.disableEndpointDiscovery
88
- }
89
- });
90
- this._container = this._client
91
- .database(this._config.databaseId)
92
- .container(this._config.containerId);
100
+ this._mutexTimeoutMs = Coerce.integer(options.config.mutexTimeoutMs);
101
+ this._instanceId = RandomHelper.generateUuidV7("compact");
102
+ }
103
+ /**
104
+ * The component needs to be stopped when the node is closed.
105
+ * @param nodeLoggingComponentType The node logging component type.
106
+ * @returns Nothing.
107
+ */
108
+ async stop(nodeLoggingComponentType) {
109
+ await ConnectionHelper.closeClient("cosmosDbClients", this.createClientId(), this._instanceId, this._mutexTimeoutMs, async (client) => client.dispose());
93
110
  }
94
111
  /**
95
112
  * Initialize the Cosmos DB environment.
@@ -98,6 +115,7 @@ export class CosmosDbEntityStorageConnector {
98
115
  */
99
116
  async bootstrap(nodeLoggingComponentType) {
100
117
  const nodeLogging = ComponentFactory.getIfExists(nodeLoggingComponentType);
118
+ const client = await this.getClient();
101
119
  // Create the database if it does not exist
102
120
  try {
103
121
  const databaseExists = await this.databaseExists();
@@ -122,7 +140,7 @@ export class CosmosDbEntityStorageConnector {
122
140
  databaseId: this._config.databaseId
123
141
  }
124
142
  });
125
- await this._client.databases.create({
143
+ await client.databases.create({
126
144
  id: this._config.databaseId
127
145
  });
128
146
  await this.waitForDatabaseExists();
@@ -165,12 +183,13 @@ export class CosmosDbEntityStorageConnector {
165
183
  containerId: this._config.containerId
166
184
  }
167
185
  });
168
- await this._client.database(this._config.databaseId).containers.create({
186
+ await client.database(this._config.databaseId).containers.create({
169
187
  id: this._config.containerId,
170
188
  partitionKey: {
171
189
  kind: PartitionKeyKind.Hash,
172
190
  paths: [`/${CosmosDbEntityStorageConnector._PARTITION_KEY}`]
173
- }
191
+ },
192
+ indexingPolicy: this.buildIndexingPolicy()
174
193
  }, { offerThroughput: this._config.offerThroughput });
175
194
  await this.waitForContainerExists();
176
195
  }
@@ -203,13 +222,12 @@ export class CosmosDbEntityStorageConnector {
203
222
  */
204
223
  async health() {
205
224
  try {
206
- await this._client
207
- .database(this._config.databaseId)
208
- .container(this._config.containerId)
209
- .read();
225
+ const container = await this.getContainer();
226
+ await container.read();
210
227
  return [
211
228
  {
212
229
  source: CosmosDbEntityStorageConnector.CLASS_NAME,
230
+ category: HealthCategory.Connectivity,
213
231
  status: HealthStatus.Ok,
214
232
  description: "healthDescription",
215
233
  data: { databaseId: this._config.databaseId, containerId: this._config.containerId }
@@ -220,6 +238,7 @@ export class CosmosDbEntityStorageConnector {
220
238
  return [
221
239
  {
222
240
  source: CosmosDbEntityStorageConnector.CLASS_NAME,
241
+ category: HealthCategory.Connectivity,
223
242
  status: HealthStatus.Error,
224
243
  description: "healthDescription",
225
244
  message: "connectionFailed",
@@ -248,11 +267,15 @@ export class CosmosDbEntityStorageConnector {
248
267
  const contextIds = await ContextIdStore.getContextIds();
249
268
  const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
250
269
  try {
270
+ const container = await this.getContainer();
251
271
  // No secondary index or conditions
252
272
  if (Is.empty(secondaryIndex) && !Is.arrayValue(conditions)) {
253
- const { resource: item } = await this._container
273
+ const { resource: item } = await container
254
274
  .item(id, partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE)
255
275
  .read();
276
+ if (Is.empty(item)) {
277
+ return undefined;
278
+ }
256
279
  return this.itemToEntity(item);
257
280
  }
258
281
  const conditionValues = [];
@@ -286,7 +309,7 @@ export class CosmosDbEntityStorageConnector {
286
309
  query: `SELECT * FROM c WHERE ${whereQuery.join(" AND ")}`,
287
310
  parameters: conditionValues
288
311
  };
289
- const { resources: items } = await this._container.items.query(query).fetchAll();
312
+ const { resources: items } = await container.items.query(query).fetchAll();
290
313
  if (items.length === 1) {
291
314
  return this.itemToEntity(items[0]);
292
315
  }
@@ -314,25 +337,66 @@ export class CosmosDbEntityStorageConnector {
314
337
  EntityStorageHelper.validateConditions(this._entitySchema, conditions);
315
338
  const contextIds = await ContextIdStore.getContextIds();
316
339
  const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
340
+ const submittedVersion = Is.stringValue(this._versionKey)
341
+ ? ObjectHelper.propertyGet(entity, this._versionKey)
342
+ : undefined;
343
+ const hasVersionCheck = !Is.empty(this._versionKey) && !Is.empty(submittedVersion) && submittedVersion > 0;
317
344
  const prepared = EntityStorageHelper.prepareEntity(entity, this._entitySchema, undefined, {
318
345
  nullBehavior: "omit"
319
346
  });
320
347
  const id = prepared[this._primaryKey.property];
348
+ const optimisticMutexKey = Is.stringValue(this._versionKey)
349
+ ? this.buildOptimisticMutexKey(partitionKey, id)
350
+ : undefined;
351
+ if (Is.stringValue(optimisticMutexKey)) {
352
+ await Mutex.lock(optimisticMutexKey, {
353
+ throwOnTimeout: true,
354
+ timeoutMs: this._mutexTimeoutMs
355
+ });
356
+ }
321
357
  try {
322
- if (Is.arrayValue(conditions)) {
323
- const item = this._container.item(id, partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE);
358
+ const container = await this.getContainer();
359
+ let itemEtag;
360
+ if (Is.stringValue(this._versionKey) || Is.arrayValue(conditions)) {
361
+ const item = container.item(id, partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE);
324
362
  const { resource: itemData } = await item.read();
325
- if (Is.notEmpty(itemData) && !this.verifyConditions(conditions, itemData)) {
326
- return;
363
+ if (!Is.empty(itemData)) {
364
+ if (hasVersionCheck) {
365
+ const storedVersion = ObjectHelper.propertyGet(itemData, this._versionKey) ?? 0;
366
+ if (storedVersion !== submittedVersion) {
367
+ throw new ConflictError(CosmosDbEntityStorageConnector.CLASS_NAME, "optimisticLockFailed", id);
368
+ }
369
+ }
370
+ if (Is.arrayValue(conditions) && !this.verifyConditions(conditions, itemData)) {
371
+ if (Is.stringValue(this._versionKey)) {
372
+ throw new ConflictError(CosmosDbEntityStorageConnector.CLASS_NAME, "conditionFailed", id);
373
+ }
374
+ return;
375
+ }
376
+ itemEtag = ObjectHelper.propertyGet(itemData, "_etag");
377
+ }
378
+ if (Is.stringValue(this._versionKey)) {
379
+ const storedVersion = !Is.empty(itemData)
380
+ ? (ObjectHelper.propertyGet(itemData, this._versionKey) ?? 0)
381
+ : 0;
382
+ ObjectHelper.propertySet(prepared, this._versionKey, storedVersion + 1);
327
383
  }
328
384
  }
329
- await this._container.items.upsert({
385
+ await container.items.upsert({
330
386
  id,
331
387
  [CosmosDbEntityStorageConnector._PARTITION_KEY]: partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE,
332
388
  ...prepared
333
- });
389
+ }, Is.stringValue(itemEtag)
390
+ ? { accessCondition: { type: "IfMatch", condition: itemEtag } }
391
+ : undefined);
334
392
  }
335
393
  catch (err) {
394
+ if (BaseError.isErrorName(err, ConflictError.CLASS_NAME)) {
395
+ throw err;
396
+ }
397
+ if (Is.object(err) && err.code === 412) {
398
+ throw new ConflictError(CosmosDbEntityStorageConnector.CLASS_NAME, "optimisticLockFailed", id);
399
+ }
336
400
  if (BaseError.isAggregateError(err)) {
337
401
  const errors = BaseError.fromAggregate(err);
338
402
  if (BaseError.someErrorCode(errors, "ResourceNotFoundException")) {
@@ -345,6 +409,11 @@ export class CosmosDbEntityStorageConnector {
345
409
  id
346
410
  }, err);
347
411
  }
412
+ finally {
413
+ if (Is.stringValue(optimisticMutexKey)) {
414
+ Mutex.unlock(optimisticMutexKey);
415
+ }
416
+ }
348
417
  }
349
418
  /**
350
419
  * Set multiple entities in a batch.
@@ -359,15 +428,30 @@ export class CosmosDbEntityStorageConnector {
359
428
  nullBehavior: "omit"
360
429
  }));
361
430
  try {
362
- await this._container.items.executeBulkOperations(preparedEntities.map(prepared => ({
363
- operationType: BulkOperationType.Upsert,
364
- partitionKey: partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE,
365
- resourceBody: {
366
- id: prepared[this._primaryKey.property],
367
- [CosmosDbEntityStorageConnector._PARTITION_KEY]: partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE,
368
- ...prepared
431
+ const container = await this.getContainer();
432
+ const pk = partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE;
433
+ const chunkSize = CosmosDbEntityStorageConnector._BATCH_CHUNK_SIZE;
434
+ const concurrency = CosmosDbEntityStorageConnector._WRITE_CONCURRENCY;
435
+ const windowSize = chunkSize * concurrency;
436
+ for (let offset = 0; offset < preparedEntities.length; offset += windowSize) {
437
+ const window = preparedEntities.slice(offset, offset + windowSize);
438
+ const sends = [];
439
+ for (let j = 0; j < window.length; j += chunkSize) {
440
+ const chunk = window.slice(j, j + chunkSize);
441
+ sends.push((async () => {
442
+ await container.items.executeBulkOperations(chunk.map(prepared => ({
443
+ operationType: BulkOperationType.Upsert,
444
+ partitionKey: pk,
445
+ resourceBody: {
446
+ id: prepared[this._primaryKey.property],
447
+ [CosmosDbEntityStorageConnector._PARTITION_KEY]: pk,
448
+ ...prepared
449
+ }
450
+ })));
451
+ })());
369
452
  }
370
- })));
453
+ await Promise.all(sends);
454
+ }
371
455
  }
372
456
  catch (err) {
373
457
  throw new GeneralError(CosmosDbEntityStorageConnector.CLASS_NAME, "setBatchFailed", undefined, err);
@@ -382,10 +466,11 @@ export class CosmosDbEntityStorageConnector {
382
466
  const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
383
467
  const pk = partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE;
384
468
  try {
469
+ const container = await this.getContainer();
385
470
  let continuationToken;
386
471
  do {
387
472
  const feedOptions = { maxItemCount: 100, continuationToken };
388
- const { resources, continuationToken: nextToken } = await this._container.items
473
+ const { resources, continuationToken: nextToken } = await container.items
389
474
  .query({
390
475
  query: `SELECT c.id FROM c WHERE c.${CosmosDbEntityStorageConnector._PARTITION_KEY} = @pk`,
391
476
  parameters: [{ name: "@pk", value: pk }]
@@ -398,7 +483,7 @@ export class CosmosDbEntityStorageConnector {
398
483
  id: r.id,
399
484
  partitionKey: pk
400
485
  }));
401
- await this._container.items.executeBulkOperations(operations);
486
+ await container.items.executeBulkOperations(operations);
402
487
  }
403
488
  } while (Is.stringValue(continuationToken));
404
489
  }
@@ -417,17 +502,33 @@ export class CosmosDbEntityStorageConnector {
417
502
  EntityStorageHelper.validateConditions(this._entitySchema, conditions);
418
503
  const contextIds = await ContextIdStore.getContextIds();
419
504
  const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
505
+ const optimisticMutexKey = Is.stringValue(this._versionKey)
506
+ ? this.buildOptimisticMutexKey(partitionKey, id)
507
+ : undefined;
508
+ if (Is.stringValue(optimisticMutexKey)) {
509
+ await Mutex.lock(optimisticMutexKey, {
510
+ throwOnTimeout: true,
511
+ timeoutMs: this._mutexTimeoutMs
512
+ });
513
+ }
420
514
  try {
421
- const item = this._container.item(id, partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE);
515
+ const container = await this.getContainer();
516
+ const item = container.item(id, partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE);
422
517
  const { resource: itemData } = await item.read();
423
518
  if (Is.notEmpty(itemData)) {
424
519
  if (Is.arrayValue(conditions) && !this.verifyConditions(conditions, itemData)) {
520
+ if (Is.stringValue(this._versionKey)) {
521
+ throw new ConflictError(CosmosDbEntityStorageConnector.CLASS_NAME, "conditionFailed", id);
522
+ }
425
523
  return;
426
524
  }
427
525
  await item.delete();
428
526
  }
429
527
  }
430
528
  catch (err) {
529
+ if (BaseError.isErrorName(err, ConflictError.CLASS_NAME)) {
530
+ throw err;
531
+ }
431
532
  if (BaseError.fromError(err) &&
432
533
  Is.object(err) &&
433
534
  err.body?.code === "NotFound") {
@@ -437,6 +538,11 @@ export class CosmosDbEntityStorageConnector {
437
538
  id
438
539
  }, err);
439
540
  }
541
+ finally {
542
+ if (Is.stringValue(optimisticMutexKey)) {
543
+ Mutex.unlock(optimisticMutexKey);
544
+ }
545
+ }
440
546
  }
441
547
  /**
442
548
  * Remove multiple entities by id.
@@ -448,12 +554,13 @@ export class CosmosDbEntityStorageConnector {
448
554
  const contextIds = await ContextIdStore.getContextIds();
449
555
  const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
450
556
  try {
557
+ const container = await this.getContainer();
451
558
  const operations = ids.map(id => ({
452
559
  operationType: BulkOperationType.Delete,
453
560
  id,
454
561
  partitionKey: partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE
455
562
  }));
456
- await this._container.items.executeBulkOperations(operations);
563
+ await container.items.executeBulkOperations(operations);
457
564
  }
458
565
  catch (err) {
459
566
  throw new GeneralError(CosmosDbEntityStorageConnector.CLASS_NAME, "removeBatchFailed", undefined, err);
@@ -474,8 +581,9 @@ export class CosmosDbEntityStorageConnector {
474
581
  data: { containerId: this._config.containerId }
475
582
  });
476
583
  try {
584
+ const container = await this.getContainer();
477
585
  if (await this.containerExists()) {
478
- await this._container.delete();
586
+ await container.delete();
479
587
  await this.waitForContainerNotExists();
480
588
  }
481
589
  await nodeLogging?.log({
@@ -515,22 +623,35 @@ export class CosmosDbEntityStorageConnector {
515
623
  EntityStorageHelper.validateSortProperties(this._entitySchema, sortProperties);
516
624
  EntityStorageHelper.validateProperties(this._entitySchema, properties);
517
625
  EntityStorageHelper.validateConditionProperties(this._entitySchema, conditions);
626
+ // Only the sort shapes covered by the composite indexes from buildIndexingPolicy are accepted.
627
+ if (Is.arrayValue(sortProperties)) {
628
+ const nonPrimarySorts = sortProperties.filter(sortProperty => sortProperty.property !== this._primaryKey.property);
629
+ if (nonPrimarySorts.length > 1) {
630
+ throw new GeneralError(CosmosDbEntityStorageConnector.CLASS_NAME, "sortUnsupported", {
631
+ properties: sortProperties.map(sortProperty => sortProperty.property),
632
+ primaryKey: this._primaryKey.property
633
+ });
634
+ }
635
+ }
518
636
  if (!Is.empty(limit)) {
519
637
  const validationFailures = [];
520
638
  Validation.integer("limit", limit, validationFailures, undefined, { minValue: 1 });
521
639
  Validation.asValidationError(CosmosDbEntityStorageConnector.CLASS_NAME, "query", validationFailures);
522
640
  }
523
641
  try {
642
+ const container = await this.getContainer();
524
643
  const returnSize = limit ?? CosmosDbEntityStorageConnector._DEFAULT_LIMIT;
525
644
  let orderByClause = "";
526
- if (Array.isArray(sortProperties)) {
527
- if (sortProperties.length > 1) {
528
- throw new GeneralError(CosmosDbEntityStorageConnector.CLASS_NAME, "sortSingle");
529
- }
530
- for (const sortProperty of sortProperties) {
645
+ if (Is.arrayValue(sortProperties)) {
646
+ // The primary key is unique so no property after its first occurrence can affect
647
+ // the order, truncating keeps the ORDER BY within the provisioned composite indexes.
648
+ const primaryKeyIndex = sortProperties.findIndex(sortProperty => sortProperty.property === this._primaryKey.property);
649
+ const effectiveSorts = primaryKeyIndex === -1 ? sortProperties : sortProperties.slice(0, primaryKeyIndex + 1);
650
+ const orderClauses = effectiveSorts.map(sortProperty => {
531
651
  const direction = sortProperty.sortDirection === SortDirection.Ascending ? "asc" : "desc";
532
- orderByClause = `ORDER BY c.${String(sortProperty.property)} ${direction}`;
533
- }
652
+ return `c.${String(sortProperty.property)} ${direction}`;
653
+ });
654
+ orderByClause = `ORDER BY ${orderClauses.join(", ")}`;
534
655
  }
535
656
  const attributeNames = {};
536
657
  const attributeValues = {};
@@ -548,13 +669,13 @@ export class CosmosDbEntityStorageConnector {
548
669
  ...Object.keys(attributeValues).map(key => ({ name: `@${key}`, value: attributeValues[key] }))
549
670
  ];
550
671
  const queryPartitionKey = partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE;
551
- // For sorted queries use OFFSET LIMIT — CosmosDB continuation tokens are
672
+ // For sorted queries use OFFSET LIMIT - CosmosDB continuation tokens are
552
673
  // not reliably returned for ORDER BY queries across all service versions.
553
674
  // The cursor is a numeric offset encoded as a string (same as SQL connectors).
554
675
  const startIndex = Coerce.number(cursor) ?? 0;
555
676
  sql = `${baseQuery} ${orderByClause} OFFSET ${startIndex} LIMIT ${returnSize + 1}`;
556
677
  const sortedQuerySpecs = { query: sql, parameters: queryParameters };
557
- const sortedResponse = await this._container.items
678
+ const sortedResponse = await container.items
558
679
  .query(sortedQuerySpecs, {
559
680
  partitionKey: queryPartitionKey,
560
681
  maxItemCount: returnSize + 1
@@ -579,6 +700,7 @@ export class CosmosDbEntityStorageConnector {
579
700
  async count(conditions) {
580
701
  EntityStorageHelper.validateConditionProperties(this._entitySchema, conditions);
581
702
  try {
703
+ const container = await this.getContainer();
582
704
  const contextIds = await ContextIdStore.getContextIds();
583
705
  const partitionKey = ContextIdHelper.combinedContextKey(contextIds, this._partitionContextIds);
584
706
  const attributeNames = {};
@@ -599,7 +721,7 @@ export class CosmosDbEntityStorageConnector {
599
721
  ...Object.keys(attributeValues).map(key => ({ name: `@${key}`, value: attributeValues[key] }))
600
722
  ]
601
723
  };
602
- const { resources } = await this._container.items.query(querySpec).fetchAll();
724
+ const { resources } = await container.items.query(querySpec).fetchAll();
603
725
  return resources[0] ?? 0;
604
726
  }
605
727
  catch (err) {
@@ -608,27 +730,58 @@ export class CosmosDbEntityStorageConnector {
608
730
  }
609
731
  /**
610
732
  * Get a unique list of all the context ids from the storage.
733
+ * @param loggingComponentType The optional component type to use for logging skipped partition ids.
611
734
  * @returns The list of unique context ids.
612
735
  */
613
- async getPartitionContextIds() {
736
+ async getPartitionContextIds(loggingComponentType) {
614
737
  const partitionContextIds = this._partitionContextIds;
615
738
  if (!Is.arrayValue(partitionContextIds)) {
616
739
  return undefined;
617
740
  }
618
741
  try {
619
- const { resources: partitionIds } = await this._container.items
742
+ const container = await this.getContainer();
743
+ const { resources: partitionIds } = await container.items
620
744
  .query({
621
745
  query: `SELECT DISTINCT VALUE c.${CosmosDbEntityStorageConnector._PARTITION_KEY} FROM c`
622
746
  })
623
747
  .fetchAll();
624
- return partitionIds
625
- .filter(id => Is.stringValue(id))
626
- .map(id => ContextIdHelper.shortSplit(partitionContextIds, id));
748
+ const contextIds = [];
749
+ const skipped = [];
750
+ for (const partitionId of partitionIds.filter(id => Is.stringValue(id))) {
751
+ const split = EntityStorageHelper.tryShortSplit(partitionContextIds, partitionId);
752
+ if (Is.undefined(split)) {
753
+ skipped.push(partitionId);
754
+ }
755
+ else {
756
+ contextIds.push(split);
757
+ }
758
+ }
759
+ if (Is.arrayValue(skipped)) {
760
+ const nodeLogging = ComponentFactory.getIfExists(loggingComponentType);
761
+ await nodeLogging?.log({
762
+ level: "warn",
763
+ source: CosmosDbEntityStorageConnector.CLASS_NAME,
764
+ ts: Date.now(),
765
+ message: "partitionIdsSkipped",
766
+ data: {
767
+ expected: partitionContextIds.length,
768
+ partitionIds: skipped.join(", ")
769
+ }
770
+ });
771
+ }
772
+ return contextIds;
627
773
  }
628
774
  catch (err) {
629
775
  throw new GeneralError(CosmosDbEntityStorageConnector.CLASS_NAME, "getPartitionContextIdsFailed", undefined, err);
630
776
  }
631
777
  }
778
+ /**
779
+ * Get the connector implementation version.
780
+ * @returns The connector implementation version.
781
+ */
782
+ connectorVersion() {
783
+ return 0;
784
+ }
632
785
  /**
633
786
  * Create the target connector for performing the migration using a temporary container.
634
787
  * @param newEntitySchema The name of the new entity schema to create the connector for.
@@ -706,7 +859,7 @@ export class CosmosDbEntityStorageConnector {
706
859
  let continuationToken;
707
860
  do {
708
861
  const feedOptions = { maxItemCount: batchSize, continuationToken };
709
- const { resources, continuationToken: nextToken } = await sourceConnector._container.items
862
+ const { resources, continuationToken: nextToken } = await (await sourceConnector.getContainer()).items
710
863
  .query({
711
864
  query: `SELECT * FROM c WHERE c.${CosmosDbEntityStorageConnector._PARTITION_KEY} = @partitionId`,
712
865
  parameters: [{ name: "@partitionId", value: partitionKey }]
@@ -719,7 +872,7 @@ export class CosmosDbEntityStorageConnector {
719
872
  partitionKey,
720
873
  resourceBody: rest
721
874
  }));
722
- await destConnector._container.items.executeBulkOperations(operations);
875
+ await (await destConnector.getContainer()).items.executeBulkOperations(operations);
723
876
  }
724
877
  } while (Is.stringValue(continuationToken));
725
878
  }
@@ -788,7 +941,7 @@ export class CosmosDbEntityStorageConnector {
788
941
  else if (Is.array(comparator.value)) {
789
942
  const dbValues = comparator.value.map(v => this.propertyToDbValue(v, type));
790
943
  if (dbValues.length === 0 && comparator.comparison === ComparisonOperator.In) {
791
- // CosmosDB rejects `IN ()` — return always-false sentinel (#141).
944
+ // CosmosDB rejects `IN ()` - return always-false sentinel (#141).
792
945
  return "1=0";
793
946
  }
794
947
  const arrAttributeNames = [];
@@ -947,6 +1100,16 @@ export class CosmosDbEntityStorageConnector {
947
1100
  verifyConditions(conditions, obj) {
948
1101
  return conditions.every(condition => ObjectHelper.propertyGet(obj, condition.property) === condition.value);
949
1102
  }
1103
+ /**
1104
+ * Build a mutex key for optimistic-locking critical sections.
1105
+ * @param partitionKey The resolved partition key.
1106
+ * @param id The entity id.
1107
+ * @returns The mutex key.
1108
+ * @internal
1109
+ */
1110
+ buildOptimisticMutexKey(partitionKey, id) {
1111
+ return `${CosmosDbEntityStorageConnector.CLASS_NAME}:optimistic:${this._config.databaseId}:${this._config.containerId}:${partitionKey ?? CosmosDbEntityStorageConnector._PARTITION_KEY_VALUE}:${id}`;
1112
+ }
950
1113
  /**
951
1114
  * Convert an entity to an item.
952
1115
  * @param item The item to convert.
@@ -963,6 +1126,38 @@ export class CosmosDbEntityStorageConnector {
963
1126
  "_ts"
964
1127
  ]);
965
1128
  }
1129
+ /**
1130
+ * Retrieve (or lazily create) the shared Cosmos DB client for this endpoint.
1131
+ * @returns The shared client.
1132
+ * @internal
1133
+ */
1134
+ async getClient() {
1135
+ return ConnectionHelper.openClient("cosmosDbClients", this.createClientId(), this._instanceId, this._mutexTimeoutMs, async () => new CosmosClient({
1136
+ endpoint: this._config.endpoint,
1137
+ key: this._config.key,
1138
+ connectionPolicy: {
1139
+ enableEndpointDiscovery: !this._config.disableEndpointDiscovery
1140
+ }
1141
+ }));
1142
+ }
1143
+ /**
1144
+ * Get the container for this connector's configured database and container.
1145
+ * @returns The container reference.
1146
+ * @internal
1147
+ */
1148
+ async getContainer() {
1149
+ return (await this.getClient())
1150
+ .database(this._config.databaseId)
1151
+ .container(this._config.containerId);
1152
+ }
1153
+ /**
1154
+ * Build a stable cache key for the shared client based on connection parameters.
1155
+ * @returns The client cache key.
1156
+ * @internal
1157
+ */
1158
+ createClientId() {
1159
+ return `${this._config.endpoint}|${this._config.databaseId}|${this._config.containerId}`;
1160
+ }
966
1161
  /**
967
1162
  * Check if the database exists.
968
1163
  * @returns True if the database exists, false otherwise.
@@ -970,7 +1165,8 @@ export class CosmosDbEntityStorageConnector {
970
1165
  */
971
1166
  async databaseExists() {
972
1167
  try {
973
- const { resources: databaseList } = await this._client.databases.readAll().fetchAll();
1168
+ const client = await this.getClient();
1169
+ const { resources: databaseList } = await client.databases.readAll().fetchAll();
974
1170
  return databaseList.some((db) => db.id === this._config.databaseId);
975
1171
  }
976
1172
  catch {
@@ -991,6 +1187,41 @@ export class CosmosDbEntityStorageConnector {
991
1187
  await new Promise(resolve => setTimeout(resolve, 250));
992
1188
  }
993
1189
  }
1190
+ /**
1191
+ * Build the composite indexes needed to serve multi-property ORDER BY queries.
1192
+ * Pairing each sortable property with the primary key in both directions covers all
1193
+ * four direction combinations, as Cosmos DB also serves each index reversed.
1194
+ * @returns The indexing policy for the container, or undefined if the schema has no
1195
+ * sortable properties.
1196
+ * @internal
1197
+ */
1198
+ buildIndexingPolicy() {
1199
+ const primaryKeyPath = `/${this._primaryKey.property}`;
1200
+ const compositeIndexes = [];
1201
+ if (Is.arrayValue(this._entitySchema.properties)) {
1202
+ for (const prop of this._entitySchema.properties) {
1203
+ if (!prop.isPrimary && (Is.stringValue(prop.sortDirection) || prop.isSecondary)) {
1204
+ const propertyPath = `/${prop.property}`;
1205
+ compositeIndexes.push([
1206
+ { path: propertyPath, order: "ascending" },
1207
+ { path: primaryKeyPath, order: "ascending" }
1208
+ ]);
1209
+ compositeIndexes.push([
1210
+ { path: propertyPath, order: "ascending" },
1211
+ { path: primaryKeyPath, order: "descending" }
1212
+ ]);
1213
+ }
1214
+ }
1215
+ }
1216
+ return compositeIndexes.length > 0
1217
+ ? {
1218
+ indexingMode: "consistent",
1219
+ automatic: true,
1220
+ includedPaths: [{ path: "/*" }],
1221
+ compositeIndexes
1222
+ }
1223
+ : undefined;
1224
+ }
994
1225
  /**
995
1226
  * Check if the container exists.
996
1227
  * @returns True if the container exists, false otherwise.
@@ -998,7 +1229,8 @@ export class CosmosDbEntityStorageConnector {
998
1229
  */
999
1230
  async containerExists() {
1000
1231
  try {
1001
- const { resources: containers } = await this._client
1232
+ const client = await this.getClient();
1233
+ const { resources: containers } = await client
1002
1234
  .database(this._config.databaseId)
1003
1235
  .containers.readAll()
1004
1236
  .fetchAll();