feat(cosmos): add RU consumption logs (#204)

Author: PMerlet

packages/datasource-cosmos/src/index.ts

@@ -4,6 +4,7 @@ import { DataSourceFactory, Logger } from '@forestadmin/datasource-toolkit';
 import CosmosDataSource from './datasource';
 import { ConfigurationOptions, CosmosDatasourceBuilder } from './introspection/builder';
 import Introspector, { VirtualArrayCollectionConfig } from './introspection/introspector';
+import { configureRuLogging } from './model-builder/model';
 import { ManualSchemaConfig } from './types/manual-schema';
 import { convertManualSchemaToModels } from './utils/manual-schema-converter';
 import VirtualCollectionManager from './virtual-collection-manager';
@@ -59,6 +60,7 @@ export { default as TypeConverter } from './utils/type-converter';
 export { default as ModelCosmos } from './model-builder/model';
 export { configurePaginationCache, getSharedPaginationCache } from './model-builder/model';
 export { configureRetryOptions, getSharedRetryOptions } from './model-builder/model';
+export { configureRuLogging, isRuLoggingEnabled } from './model-builder/model';
 export { default as PaginationCache } from './utils/pagination-cache';
 export {
   withRetry,
@@ -169,6 +171,12 @@ export function createCosmosDataSource(
      * Allows you to define collections and their fields explicitly
      */
     schema?: ManualSchemaConfig;
+    /**
+     * Enable logging of Request Unit (RU) consumption for each Cosmos DB operation
+     * Useful for monitoring and optimizing database costs
+     * Default: false
+     */
+    logRuConsumption?: boolean;
   },
 ): DataSourceFactory {
   return async (logger: Logger) => {
@@ -186,8 +194,14 @@ export function createCosmosDataSource(
       introspectionConfig,
       disableIntrospection,
       schema,
+      logRuConsumption,
     } = options || {};
 
+    // Configure RU logging if enabled
+    if (logRuConsumption) {
+      configureRuLogging(true, logger);
+    }
+
     // Apply introspection config defaults
     const finalIntrospectionConfig: IntrospectionConfig = {
       sampleSize: introspectionConfig?.sampleSize ?? 100,
@@ -281,6 +295,12 @@ export function createCosmosDataSourceForEmulator(
     introspectionConfig?: IntrospectionConfig;
     disableIntrospection?: boolean;
     schema?: ManualSchemaConfig;
+    /**
+     * Enable logging of Request Unit (RU) consumption for each Cosmos DB operation
+     * Useful for monitoring and optimizing database costs
+     * Default: false
+     */
+    logRuConsumption?: boolean;
   },
 ): DataSourceFactory {
   // Default emulator connection details

packages/datasource-cosmos/src/model-builder/model.ts

@@ -1,5 +1,5 @@
 import { Container, CosmosClient, ItemDefinition, SqlQuerySpec } from '@azure/cosmos';
-import { RecordData } from '@forestadmin/datasource-toolkit';
+import { Logger, RecordData } from '@forestadmin/datasource-toolkit';
 import { randomUUID } from 'crypto';
 
 import { OverrideTypeConverter } from '../introspection/builder';
@@ -14,6 +14,29 @@ import Serializer from '../utils/serializer';
 
 export { RetryOptions, configureRetryOptions, getSharedRetryOptions };
 
+/**
+ * Shared RU logging configuration
+ */
+let sharedRuLoggingEnabled = false;
+let sharedRuLogger: Logger | null = null;
+
+/**
+ * Configure RU logging for all Cosmos DB operations
+ * @param enabled Whether to enable RU consumption logging
+ * @param logger The Forest Admin logger to use for RU logging
+ */
+export function configureRuLogging(enabled: boolean, logger?: Logger): void {
+  sharedRuLoggingEnabled = enabled;
+  sharedRuLogger = logger ?? null;
+}
+
+/**
+ * Check if RU logging is enabled
+ */
+export function isRuLoggingEnabled(): boolean {
+  return sharedRuLoggingEnabled;
+}
+
 export interface CosmosSchema {
   [key: string]: {
     type: string;
@@ -154,12 +177,13 @@ export default class ModelCosmos {
       };
 
       // eslint-disable-next-line no-await-in-loop -- Sequential execution maintains order
-      const { resource } = await withRetry(
+      const response = await withRetry(
         () => this.container.items.create(itemToCreate),
         this.retryOptions,
       );
+      this.logRuConsumption('create', response.requestCharge);
       // Flatten and serialize the response
-      createdRecords.push(Serializer.serialize(resource));
+      createdRecords.push(Serializer.serialize(response.resource));
     }
 
     return createdRecords;
@@ -174,21 +198,23 @@ export default class ModelCosmos {
     for (const { id, partitionKey } of items) {
       // Point read: directly fetch item using id + partition key (1 RU)
       // eslint-disable-next-line no-await-in-loop -- Sequential maintains consistency
-      const { resource: existingItem } = await withRetry(
+      const readResponse = await withRetry(
         () => this.container.item(id, partitionKey).read(),
         this.retryOptions,
       );
+      this.logRuConsumption('update.read', readResponse.requestCharge);
 
-      if (existingItem) {
+      if (readResponse.resource) {
         // Deep merge the unflattened patch with the existing item
-        const updatedItem = Serializer.deepMerge(existingItem, unflattenedPatch);
+        const updatedItem = Serializer.deepMerge(readResponse.resource, unflattenedPatch);
 
         // Point update: directly update using id + partition key
         // eslint-disable-next-line no-await-in-loop -- Sequential maintains consistency
-        await withRetry(
+        const replaceResponse = await withRetry(
           () => this.container.item(id, partitionKey).replace(updatedItem),
           this.retryOptions,
         );
+        this.logRuConsumption('update.replace', replaceResponse.requestCharge);
       }
     }
   }
@@ -198,7 +224,11 @@ export default class ModelCosmos {
     for (const { id, partitionKey } of items) {
       // Point delete: directly delete using id + partition key
       // eslint-disable-next-line no-await-in-loop -- Sequential maintains consistency
-      await withRetry(() => this.container.item(id, partitionKey).delete(), this.retryOptions);
+      const response = await withRetry(
+        () => this.container.item(id, partitionKey).delete(),
+        this.retryOptions,
+      );
+      this.logRuConsumption('delete', response.requestCharge);
     }
   }
 
@@ -222,12 +252,13 @@ export default class ModelCosmos {
 
     // If no pagination parameters, fetch all (backward compatibility)
     if (offset === undefined && limit === undefined) {
-      const { resources } = await withRetry(
+      const response = await withRetry(
         () => this.container.items.query(querySpec, queryOptions).fetchAll(),
         this.retryOptions,
       );
+      this.logRuConsumption('query.fetchAll', response.requestCharge);
 
-      return resources.map(item => Serializer.serialize(item));
+      return response.resources.map(item => Serializer.serialize(item));
     }
 
     const effectiveOffset = offset || 0;
@@ -281,7 +312,8 @@ export default class ModelCosmos {
     // Iterate through pages until we have enough items
     // eslint-disable-next-line no-restricted-syntax
     for await (const response of queryIterator.getAsyncIterator()) {
-      const { resources: page, continuationToken } = response;
+      const { resources: page, continuationToken, requestCharge } = response;
+      this.logRuConsumption('query.page', requestCharge);
 
       // Store continuation token for future use at regular intervals
       const cacheInterval = this.paginationCache.getCacheInterval();
@@ -336,12 +368,13 @@ export default class ModelCosmos {
       queryOptions.partitionKey = partitionKey;
     }
 
-    const { resources } = await withRetry(
+    const response = await withRetry(
       () => this.container.items.query(querySpec, queryOptions).fetchAll(),
       this.retryOptions,
     );
+    this.logRuConsumption('aggregate', response.requestCharge);
 
-    return resources.map(item => Serializer.serialize(item));
+    return response.resources.map(item => Serializer.serialize(item));
   }
 
   public async count(querySpec?: SqlQuerySpec, partitionKey?: string | number): Promise<number> {
@@ -357,22 +390,24 @@ export default class ModelCosmos {
       const countQuery: SqlQuerySpec = {
         query: 'SELECT VALUE COUNT(1) FROM c',
       };
-      const { resources } = await withRetry(
+      const response = await withRetry(
         () => this.container.items.query<number>(countQuery, queryOptions).fetchAll(),
         this.retryOptions,
       );
+      this.logRuConsumption('count', response.requestCharge);
 
-      return resources[0] || 0;
+      return response.resources[0] || 0;
     }
 
     // Convert the query to a COUNT query to avoid fetching all records
     const countQuery = this.convertToCountQuery(querySpec);
-    const { resources } = await withRetry(
+    const response = await withRetry(
       () => this.container.items.query<number>(countQuery, queryOptions).fetchAll(),
       this.retryOptions,
     );
+    this.logRuConsumption('count', response.requestCharge);
 
-    return resources[0] || 0;
+    return response.resources[0] || 0;
   }
 
   /**
@@ -440,4 +475,16 @@ export default class ModelCosmos {
   public getContainerName(): string {
     return this.containerName;
   }
+
+  /**
+   * Log RU consumption if enabled
+   */
+  private logRuConsumption(operation: string, requestCharge: number | undefined): void {
+    if (!sharedRuLoggingEnabled || requestCharge === undefined) {
+      return;
+    }
+
+    const message = `[Cosmos RU] ${this.name}.${operation}: ${requestCharge} RUs`;
+    sharedRuLogger?.('Info', message);
+  }
 }