@utaba/deep-memory-storage-cosmosdb 0.8.1 → 0.9.1

package/dist/index.cjs

@@ -41,6 +41,8 @@ var import_deep_memory5 = require("@utaba/deep-memory");
 
 // src/CosmosDbConnection.ts
 var import_gremlin = __toESM(require("gremlin"), 1);
+var import_node_async_hooks = require("async_hooks");
+var usageScope = new import_node_async_hooks.AsyncLocalStorage();
 var CosmosDbConnection = class {
   client = null;
   config;
@@ -88,11 +90,18 @@ var CosmosDbConnection = class {
         const resultSet = await this.client.submit(query, bindings);
         const items = resultSet.toArray();
         const requestCharge = extractRequestCharge(resultSet);
+        const acc = usageScope.getStore();
+        if (acc) {
+          acc.calls++;
+          if (typeof requestCharge === "number") acc.ru += requestCharge;
+        }
         return { items, requestCharge };
       } catch (err) {
         lastError = err;
         if (isTransientError(err) && attempt < this.config.maxRetries) {
           const retryAfterMs = getRetryAfterMs(err, attempt);
+          const acc = usageScope.getStore();
+          if (acc) acc.retries++;
           await sleep(retryAfterMs);
           continue;
         }
@@ -910,6 +919,7 @@ async function exploreNeighborhood(conn, repositoryId, entityId, options) {
   for (let depth = 0; depth < options.depth; depth++) {
     const layer = {};
     const nextFrontier = /* @__PURE__ */ new Set();
+    const layerCache = /* @__PURE__ */ new Map();
     for (const currentId of frontier) {
       const rels = await getRelationshipsForEntity(conn, repositoryId, currentId, options);
       for (const rel of rels) {
@@ -923,20 +933,20 @@ async function exploreNeighborhood(conn, repositoryId, entityId, options) {
         } else {
           connectedId = rel.sourceEntityId;
         }
-        if (!visited.has(connectedId)) {
-          if (options.entityTypes && options.entityTypes.length > 0) {
-            const entity = await getEntityLight(conn, repositoryId, connectedId);
-            if (!entity || !options.entityTypes.includes(entity.entityType)) continue;
-            layer[relType].entities.push(entity);
+        if (visited.has(connectedId)) continue;
+        let entity = layerCache.get(connectedId);
+        if (entity === void 0) {
+          const loaded = await getEntityLight(conn, repositoryId, connectedId);
+          if (loaded && (!options.entityTypes || options.entityTypes.length === 0 || options.entityTypes.includes(loaded.entityType))) {
+            entity = loaded;
           } else {
-            const entity = await getEntityLight(conn, repositoryId, connectedId);
-            if (entity) {
-              layer[relType].entities.push(entity);
-            }
+            entity = null;
           }
-          visited.add(connectedId);
-          nextFrontier.add(connectedId);
+          layerCache.set(connectedId, entity);
+          if (entity) nextFrontier.add(connectedId);
         }
+        if (!entity) continue;
+        layer[relType].entities.push(entity);
         layer[relType].relationships.push(rel);
         layer[relType].total = layer[relType].entities.length;
       }
@@ -951,6 +961,9 @@ async function exploreNeighborhood(conn, repositoryId, entityId, options) {
     if (Object.keys(layer).length > 0) {
       layers.push(layer);
     }
+    for (const id of nextFrontier) {
+      visited.add(id);
+    }
     frontier = nextFrontier;
     if (frontier.size === 0) break;
   }
@@ -1328,14 +1341,17 @@ async function upsertRelationship(conn, repositoryId, rel) {
 }
 
 // src/CosmosDbProvider.ts
+var PROVIDER_NAME = "cosmosdb";
 var SCHEMA_VERSION = 1;
 var META_VERTEX_ID = "_meta:schema";
 var CosmosDbProvider = class {
   conn;
   config;
   compiler = new import_deep_memory5.GremlinCompiler();
+  reportUsage;
   constructor(config) {
     this.config = config;
+    this.reportUsage = (0, import_deep_memory5.createSafeSink)(config.reportUsage);
     this.conn = new CosmosDbConnection({
       endpoint: config.endpoint,
       key: config.key,
@@ -1346,6 +1362,39 @@ var CosmosDbProvider = class {
       rejectUnauthorized: config.rejectUnauthorized
     });
   }
+  /**
+   * Run the given operation inside an async-local usage scope. While the
+   * operation runs, every `conn.submit()` call accumulates its request charge
+   * into a shared accumulator. When the operation completes (success or
+   * failure), a single {@link OperationUsage} record is emitted to the sink.
+   *
+   * Nested public-method calls (e.g. `traverse` internally calls
+   * `getVocabulary`) are detected: the inner call joins the outer scope
+   * instead of starting a new one, so one user-visible operation produces
+   * exactly one usage record covering all its internal round-trips.
+   *
+   * If no sink is configured, this is a zero-overhead pass-through.
+   */
+  async track(operation, repositoryId, fn) {
+    if (!this.reportUsage) return fn();
+    if (usageScope.getStore()) return fn();
+    const acc = { ru: 0, calls: 0, retries: 0 };
+    try {
+      return await usageScope.run(acc, fn);
+    } finally {
+      if (acc.calls > 0) {
+        this.reportUsage({
+          provider: PROVIDER_NAME,
+          operation,
+          unit: "RU",
+          value: acc.ru,
+          ...repositoryId ? { repositoryId } : {},
+          timestamp: /* @__PURE__ */ new Date(),
+          details: { calls: acc.calls, retries: acc.retries }
+        });
+      }
+    }
+  }
   /**
    * Reject any repositoryId that is not a valid v4 UUID.
    *
@@ -1371,6 +1420,9 @@ var CosmosDbProvider = class {
     await this.conn.close();
   }
   async ensureSchema() {
+    return this.track("ensureSchema", void 0, () => this.ensureSchemaImpl());
+  }
+  async ensureSchemaImpl() {
     const restBase = this.getRestEndpoint();
     const partitionKey = this.config.partitionKey ?? "/repositoryId";
     let databaseCreated = false;
@@ -1449,72 +1501,139 @@ var CosmosDbProvider = class {
   // ─── Repository ────────────────────────────────────────────────────
   async createRepository(config) {
     this.assertValidRepositoryId(config.repositoryId);
-    return createRepository(this.conn, config);
+    return this.track(
+      "createRepository",
+      config.repositoryId,
+      () => createRepository(this.conn, config)
+    );
   }
   async getRepository(repositoryId) {
     this.assertValidRepositoryId(repositoryId);
-    return getRepository(this.conn, repositoryId);
+    return this.track(
+      "getRepository",
+      repositoryId,
+      () => getRepository(this.conn, repositoryId)
+    );
   }
   async listRepositories(filter) {
-    return listRepositories(this.conn, filter);
+    return this.track(
+      "listRepositories",
+      void 0,
+      () => listRepositories(this.conn, filter)
+    );
   }
   async updateRepository(repositoryId, updates) {
     this.assertValidRepositoryId(repositoryId);
-    return updateRepository(this.conn, repositoryId, updates);
+    return this.track(
+      "updateRepository",
+      repositoryId,
+      () => updateRepository(this.conn, repositoryId, updates)
+    );
   }
   async deleteRepository(repositoryId, onProgress) {
     this.assertValidRepositoryId(repositoryId);
-    return deleteRepository(this.conn, repositoryId, onProgress);
+    return this.track(
+      "deleteRepository",
+      repositoryId,
+      () => deleteRepository(this.conn, repositoryId, onProgress)
+    );
   }
   async deleteAllContents(repositoryId, onProgress) {
     this.assertValidRepositoryId(repositoryId);
-    return deleteAllContents(this.conn, repositoryId, onProgress);
+    return this.track(
+      "deleteAllContents",
+      repositoryId,
+      () => deleteAllContents(this.conn, repositoryId, onProgress)
+    );
   }
   async getRepositoryStats(repositoryId) {
     this.assertValidRepositoryId(repositoryId);
-    return getRepositoryStats(this.conn, repositoryId);
+    return this.track(
+      "getRepositoryStats",
+      repositoryId,
+      () => getRepositoryStats(this.conn, repositoryId)
+    );
   }
   // ─── Vocabulary ────────────────────────────────────────────────────
   async getVocabulary(repositoryId) {
     this.assertValidRepositoryId(repositoryId);
-    return getVocabulary(this.conn, repositoryId);
+    return this.track(
+      "getVocabulary",
+      repositoryId,
+      () => getVocabulary(this.conn, repositoryId)
+    );
   }
   async saveVocabulary(repositoryId, vocabulary) {
     this.assertValidRepositoryId(repositoryId);
-    return saveVocabulary(this.conn, repositoryId, vocabulary);
+    return this.track(
+      "saveVocabulary",
+      repositoryId,
+      () => saveVocabulary(this.conn, repositoryId, vocabulary)
+    );
   }
   async getVocabularyChangeLog(repositoryId, options) {
     this.assertValidRepositoryId(repositoryId);
-    return getVocabularyChangeLog(this.conn, repositoryId, options);
+    return this.track(
+      "getVocabularyChangeLog",
+      repositoryId,
+      () => getVocabularyChangeLog(this.conn, repositoryId, options)
+    );
   }
   // ─── Entities ──────────────────────────────────────────────────────
   async createEntity(repositoryId, entity) {
     this.assertValidRepositoryId(repositoryId);
-    return createEntity(this.conn, repositoryId, entity);
+    return this.track(
+      "createEntity",
+      repositoryId,
+      () => createEntity(this.conn, repositoryId, entity)
+    );
   }
   async getEntity(repositoryId, entityId) {
     this.assertValidRepositoryId(repositoryId);
-    return getEntity(this.conn, repositoryId, entityId);
+    return this.track(
+      "getEntity",
+      repositoryId,
+      () => getEntity(this.conn, repositoryId, entityId)
+    );
   }
   async getEntityBySlug(repositoryId, slug) {
     this.assertValidRepositoryId(repositoryId);
-    return getEntityBySlug(this.conn, repositoryId, slug);
+    return this.track(
+      "getEntityBySlug",
+      repositoryId,
+      () => getEntityBySlug(this.conn, repositoryId, slug)
+    );
   }
   async getEntities(repositoryId, entityIds) {
     this.assertValidRepositoryId(repositoryId);
-    return getEntities(this.conn, repositoryId, entityIds);
+    return this.track(
+      "getEntities",
+      repositoryId,
+      () => getEntities(this.conn, repositoryId, entityIds)
+    );
   }
   async updateEntity(repositoryId, entityId, updates) {
     this.assertValidRepositoryId(repositoryId);
-    return updateEntity(this.conn, repositoryId, entityId, updates);
+    return this.track(
+      "updateEntity",
+      repositoryId,
+      () => updateEntity(this.conn, repositoryId, entityId, updates)
+    );
   }
   async deleteEntity(repositoryId, entityId) {
     this.assertValidRepositoryId(repositoryId);
-    return deleteEntity(this.conn, repositoryId, entityId);
+    return this.track(
+      "deleteEntity",
+      repositoryId,
+      () => deleteEntity(this.conn, repositoryId, entityId)
+    );
   }
   async deleteEntities(repositoryId, ids) {
     if (ids.length === 0) return { deleted: [], notFound: [] };
     this.assertValidRepositoryId(repositoryId);
+    return this.track("deleteEntities", repositoryId, () => this.deleteEntitiesImpl(repositoryId, ids));
+  }
+  async deleteEntitiesImpl(repositoryId, ids) {
     const deleted = [];
     const CHUNK = 100;
     for (let i = 0; i < ids.length; i += CHUNK) {
@@ -1553,32 +1672,59 @@ var CosmosDbProvider = class {
   }
   async deleteEntitiesByType(repositoryId, entityType) {
     this.assertValidRepositoryId(repositoryId);
-    return deleteEntitiesByType(this.conn, repositoryId, entityType);
+    return this.track(
+      "deleteEntitiesByType",
+      repositoryId,
+      () => deleteEntitiesByType(this.conn, repositoryId, entityType)
+    );
   }
   async findEntities(repositoryId, query) {
     this.assertValidRepositoryId(repositoryId);
-    return findEntities(this.conn, repositoryId, query);
+    return this.track(
+      "findEntities",
+      repositoryId,
+      () => findEntities(this.conn, repositoryId, query)
+    );
   }
   // ─── Relationships ─────────────────────────────────────────────────
   async createRelationship(repositoryId, relationship) {
     this.assertValidRepositoryId(repositoryId);
-    return createRelationship(this.conn, repositoryId, relationship);
+    return this.track(
+      "createRelationship",
+      repositoryId,
+      () => createRelationship(this.conn, repositoryId, relationship)
+    );
   }
   async getRelationship(repositoryId, relationshipId) {
     this.assertValidRepositoryId(repositoryId);
-    return getRelationship(this.conn, repositoryId, relationshipId);
+    return this.track(
+      "getRelationship",
+      repositoryId,
+      () => getRelationship(this.conn, repositoryId, relationshipId)
+    );
   }
   async getEntityRelationships(repositoryId, entityId, options) {
     this.assertValidRepositoryId(repositoryId);
-    return getEntityRelationships(this.conn, repositoryId, entityId, options);
+    return this.track(
+      "getEntityRelationships",
+      repositoryId,
+      () => getEntityRelationships(this.conn, repositoryId, entityId, options)
+    );
   }
   async deleteRelationship(repositoryId, relationshipId) {
     this.assertValidRepositoryId(repositoryId);
-    return deleteRelationship(this.conn, repositoryId, relationshipId);
+    return this.track(
+      "deleteRelationship",
+      repositoryId,
+      () => deleteRelationship(this.conn, repositoryId, relationshipId)
+    );
   }
   async deleteRelationships(repositoryId, ids) {
     if (ids.length === 0) return { deleted: [], notFound: [] };
     this.assertValidRepositoryId(repositoryId);
+    return this.track("deleteRelationships", repositoryId, () => this.deleteRelationshipsImpl(repositoryId, ids));
+  }
+  async deleteRelationshipsImpl(repositoryId, ids) {
     const deleted = [];
     const CHUNK = 100;
     for (let i = 0; i < ids.length; i += CHUNK) {
@@ -1614,30 +1760,99 @@ var CosmosDbProvider = class {
   }
   async deleteRelationshipsByType(repositoryId, relationshipType) {
     this.assertValidRepositoryId(repositoryId);
-    return deleteRelationshipsByType(this.conn, repositoryId, relationshipType);
+    return this.track(
+      "deleteRelationshipsByType",
+      repositoryId,
+      () => deleteRelationshipsByType(this.conn, repositoryId, relationshipType)
+    );
   }
   // ─── Graph Traversal (StorageProvider) ─────────────────────────────
   async exploreNeighborhood(repositoryId, entityId, options) {
     this.assertValidRepositoryId(repositoryId);
-    return exploreNeighborhood(this.conn, repositoryId, entityId, options);
+    return this.track(
+      "exploreNeighborhood",
+      repositoryId,
+      () => exploreNeighborhood(this.conn, repositoryId, entityId, options)
+    );
   }
   async findPaths(repositoryId, sourceId, targetId, options) {
     this.assertValidRepositoryId(repositoryId);
-    return findPaths(this.conn, repositoryId, sourceId, targetId, options);
+    return this.track(
+      "findPaths",
+      repositoryId,
+      () => findPaths(this.conn, repositoryId, sourceId, targetId, options)
+    );
   }
   // ─── Timeline ──────────────────────────────────────────────────────
   async getTimeline(repositoryId, entityId, options) {
     this.assertValidRepositoryId(repositoryId);
-    return getTimeline(this.conn, repositoryId, entityId, options);
+    return this.track(
+      "getTimeline",
+      repositoryId,
+      () => getTimeline(this.conn, repositoryId, entityId, options)
+    );
   }
   // ─── Bulk Operations ───────────────────────────────────────────────
   exportAll(repositoryId) {
     this.assertValidRepositoryId(repositoryId);
-    return exportAll(this.conn, repositoryId);
+    return this.trackIterable("exportAll", repositoryId, exportAll(this.conn, repositoryId));
   }
   async importBulk(repositoryId, data, options) {
     this.assertValidRepositoryId(repositoryId);
-    return importBulk(this.conn, repositoryId, data, options);
+    return this.track(
+      "importBulk",
+      repositoryId,
+      () => importBulk(this.conn, repositoryId, data, options)
+    );
+  }
+  /**
+   * Wrap an AsyncIterable so every emitted chunk is generated inside the
+   * usage scope. On each iteration step, the scope aggregates charges from
+   * the next chunk's submits; when the iterator completes (or is closed), a
+   * single usage record is emitted for the whole stream.
+   */
+  trackIterable(operation, repositoryId, source) {
+    if (!this.reportUsage) return source;
+    const sink = this.reportUsage;
+    return {
+      [Symbol.asyncIterator]: () => {
+        const iter = source[Symbol.asyncIterator]();
+        const acc = { ru: 0, calls: 0, retries: 0 };
+        let emitted = false;
+        const emit = () => {
+          if (emitted) return;
+          emitted = true;
+          if (acc.calls > 0) {
+            sink({
+              provider: PROVIDER_NAME,
+              operation,
+              unit: "RU",
+              value: acc.ru,
+              repositoryId,
+              timestamp: /* @__PURE__ */ new Date(),
+              details: { calls: acc.calls, retries: acc.retries }
+            });
+          }
+        };
+        return {
+          async next() {
+            const step = await usageScope.run(acc, () => iter.next());
+            if (step.done) emit();
+            return step;
+          },
+          async return(value) {
+            emit();
+            if (iter.return) return iter.return(value);
+            return { done: true, value };
+          },
+          async throw(err) {
+            emit();
+            if (iter.throw) return iter.throw(err);
+            throw err;
+          }
+        };
+      }
+    };
   }
   // ─── GraphTraversalProvider ────────────────────────────────────────
   getCapabilities() {
@@ -1655,6 +1870,9 @@ var CosmosDbProvider = class {
   }
   async traverse(repositoryId, spec) {
     this.assertValidRepositoryId(repositoryId);
+    return this.track("traverse", repositoryId, () => this.traverseImpl(repositoryId, spec));
+  }
+  async traverseImpl(repositoryId, spec) {
     const startTime = Date.now();
     const vocabulary = await this.getVocabulary(repositoryId);
     const compiled = this.compiler.compile(spec, vocabulary);
@@ -1748,6 +1966,9 @@ var CosmosDbProvider = class {
    * repositoryId partition predicate.
    */
   async executeNativeQuery(_repositoryId, query, params) {
+    return this.track("executeNativeQuery", void 0, () => this.executeNativeQueryImpl(query, params));
+  }
+  async executeNativeQueryImpl(query, params) {
     const startTime = Date.now();
     try {
       const result = await this.conn.submit(query, params);