fix: oom for deprecated schema explorer (#7074)

Author: n1ru4l

.changeset/wild-bags-know.md

@@ -0,0 +1,6 @@
+---
+'hive': patch
+---
+
+Prevent potential resource exhaustion in the deprecated schema explorer for large
+schemas.

packages/services/api/src/modules/operations/providers/operations-manager.ts

@@ -907,15 +907,6 @@ export class OperationsManager {
       typename: string;
     } & TargetSelector,
   ) {
-    await this.session.assertPerformAction({
-      action: 'project:describe',
-      organizationId: args.organizationId,
-      params: {
-        organizationId: args.organizationId,
-        projectId: args.projectId,
-      },
-    });
-
     const loader = this.getClientNamesPerCoordinateOfTypeLoader({
       target: args.targetId,
       period: args.period,
@@ -1118,6 +1109,14 @@ export class OperationsManager {
   }: {
     period: DateRange;
   } & TargetSelector) {
+    this.logger.debug(
+      'Count coordinates of target. (organizationId=%s, projectId=%s, targetId=%s, period=%o)',
+      organization,
+      project,
+      target,
+      period,
+    );
+
     await this.session.assertPerformAction({
       action: 'project:describe',
       organizationId: organization,
@@ -1132,6 +1131,15 @@ export class OperationsManager {
       period,
     });
 
+    this.logger.debug(
+      '%d coordinates found. (organizationId=%s, projectId=%s, targetId=%s, period=%o)',
+      rows.length,
+      organization,
+      project,
+      target,
+      period,
+    );
+
     const records: {
       [coordinate: string]: {
         total: number;

packages/services/api/src/modules/operations/providers/operations-reader.ts

@@ -2140,6 +2140,8 @@ export class OperationsReader {
   }
 
   async countCoordinatesOfTarget({ target, period }: { target: string; period: DateRange }) {
+    this.logger.debug('Count coordinates of target. (targetId=%s, period=%o)', target, period);
+
     const result = await this.clickHouse.query<{
       coordinate: string;
       total: number;
@@ -2159,6 +2161,8 @@ export class OperationsReader {
       }),
     );
 
+    this.logger.debug('%d rows found. (targetId=%s, period=%o)', result.rows, target, period);
+
     return result.data.map(row => ({
       coordinate: row.coordinate,
       total: ensureNumber(row.total),

packages/services/api/src/modules/schema/resolvers/DeprecatedSchemaExplorer.ts

@@ -1,26 +1,30 @@
 import { OperationsManager } from '../../operations/providers/operations-manager';
-import { buildGraphQLTypesFromSDL, withUsedByClients } from '../utils';
+import { buildGraphQLTypesFromSDL, memo, withUsedByClients } from '../utils';
 import type { DeprecatedSchemaExplorerResolvers } from './../../../__generated__/types';
 
 export const DeprecatedSchemaExplorer: DeprecatedSchemaExplorerResolvers = {
   types: ({ sdl, supergraph, usage }, _, { injector }) => {
     const operationsManager = injector.get(OperationsManager);
 
-    async function getStats(typename: string) {
-      const stats = await operationsManager.countCoordinatesOfTarget({
-        targetId: usage.targetId,
-        organizationId: usage.organizationId,
-        projectId: usage.projectId,
-        period: usage.period,
-      });
+    // NOTE: this function is super heavy.
+    const getStats = memo(
+      async function _getStats(typename: string) {
+        const stats = await operationsManager.countCoordinatesOfTarget({
+          targetId: usage.targetId,
+          organizationId: usage.organizationId,
+          projectId: usage.projectId,
+          period: usage.period,
+        });
 
-      return withUsedByClients(stats, {
-        selector: usage,
-        period: usage.period,
-        operationsManager,
-        typename,
-      });
-    }
+        return withUsedByClients(stats, {
+          selector: usage,
+          period: usage.period,
+          operationsManager,
+          typename,
+        });
+      },
+      arg => arg,
+    );
 
     return buildGraphQLTypesFromSDL(sdl, getStats, supergraph).sort((a, b) =>
       a.entity.name.localeCompare(b.entity.name),

packages/services/api/src/modules/schema/resolvers/SchemaExplorer.ts

@@ -22,7 +22,7 @@ import {
   Kind,
 } from 'graphql';
 import { OperationsManager } from '../../operations/providers/operations-manager';
-import { withUsedByClients } from '../utils';
+import { memo, withUsedByClients } from '../utils';
 import type { SchemaExplorerResolvers } from './../../../__generated__/types';
 
 export const SchemaExplorer: SchemaExplorerResolvers = {
@@ -162,21 +162,24 @@ export const SchemaExplorer: SchemaExplorerResolvers = {
     const typeMap = schema.getTypeMap();
     const operationsManager = injector.get(OperationsManager);
 
-    async function getStats(typename: string) {
-      const stats = await operationsManager.countCoordinatesOfTarget({
-        targetId: usage.targetId,
-        organizationId: usage.organizationId,
-        projectId: usage.projectId,
-        period: usage.period,
-      });
+    const getStats = memo(
+      async function _getStats(typename: string) {
+        const stats = await operationsManager.countCoordinatesOfTarget({
+          targetId: usage.targetId,
+          organizationId: usage.organizationId,
+          projectId: usage.projectId,
+          period: usage.period,
+        });
 
-      return withUsedByClients(stats, {
-        selector: usage,
-        period: usage.period,
-        operationsManager,
-        typename,
-      });
-    }
+        return withUsedByClients(stats, {
+          selector: usage,
+          period: usage.period,
+          operationsManager,
+          typename,
+        });
+      },
+      arg => arg,
+    );
 
     for (const typename in typeMap) {
       if (typename.startsWith('__')) {

packages/services/api/src/modules/schema/resolvers/SchemaVersion.ts

@@ -1,6 +1,7 @@
 import { createPeriod, parseDateRangeInput } from '../../../shared/helpers';
 import { buildASTSchema } from '../../../shared/schema';
 import { OperationsManager } from '../../operations/providers/operations-manager';
+import { Logger } from '../../shared/providers/logger';
 import { onlyDeprecatedDocumentNode } from '../lib/deprecated-graphql';
 import { extractSuperGraphInformation } from '../lib/federation-super-graph';
 import { stripUsedSchemaCoordinatesFromDocumentNode } from '../lib/unused-graphql';
@@ -163,6 +164,16 @@ export const SchemaVersion: SchemaVersionResolvers = {
     };
   },
   deprecatedSchema: async (version, args, { injector }) => {
+    const logger = injector.get(Logger);
+
+    logger.debug(
+      'Build deprecated schema explorer. (organizationId=%s, projectId=%s, targetId=%s, schemaVersionId=%s)',
+      version.organizationId,
+      version.projectId,
+      version.targetId,
+      version.id,
+    );
+
     const [schemaAst, supergraphAst] = await Promise.all([
       injector.get(SchemaVersionHelper).getCompositeSchemaAst(version),
       injector.get(SchemaVersionHelper).getSupergraphAst(version),
@@ -178,8 +189,26 @@ export const SchemaVersion: SchemaVersionResolvers = {
       ? parseDateRangeInput(args.period.absoluteRange)
       : createPeriod('30d');
 
+    logger.debug(
+      'Start filtering full schema SDL into deprecated schema SDL. (organizationId=%s, projectId=%s, targetId=%s, schemaVersionId=%s)',
+      version.organizationId,
+      version.projectId,
+      version.targetId,
+      version.id,
+    );
+
+    const filteredSdl = onlyDeprecatedDocumentNode(schemaAst);
+
+    logger.debug(
+      'Finished filtering full schema SDL into deprecated schema SDL. (organizationId=%s, projectId=%s, targetId=%s, schemaVersionId=%s)',
+      version.organizationId,
+      version.projectId,
+      version.targetId,
+      version.id,
+    );
+
     return {
-      sdl: onlyDeprecatedDocumentNode(schemaAst),
+      sdl: filteredSdl,
       usage: {
         period,
         organizationId: version.organizationId,

packages/services/api/src/modules/schema/utils.ts

@@ -385,3 +385,20 @@ export function stringifyDefaultValue(value: unknown): string | null {
   }
   return null;
 }
+
+export function memo<R, A, K>(fn: (arg: A) => R, cacheKeyFn: (arg: A) => K): (arg: A) => R {
+  let memoizedResult: R | null = null;
+  let memoizedKey: K | null = null;
+
+  return (arg: A) => {
+    const currentKey = cacheKeyFn(arg);
+    if (memoizedKey === currentKey) {
+      return memoizedResult!;
+    }
+
+    memoizedKey = currentKey;
+    memoizedResult = fn(arg);
+
+    return memoizedResult;
+  };
+}