From 15c1412eabaa821cdfaf2589f77ce15edb9d960b Mon Sep 17 00:00:00 2001
From: Dotan Simha <dotansimha@gmail.com>
Date: Sun, 22 Dec 2024 16:48:24 +0200
Subject: [PATCH] Improve OTEL trace and add missing information about storage
 calls (#6152)

---
 .../schema/providers/schema-manager.ts        | 45 +++++++++++++++++++
 .../schema/providers/schema-version-helper.ts | 30 +++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/packages/services/api/src/modules/schema/providers/schema-manager.ts b/packages/services/api/src/modules/schema/providers/schema-manager.ts
index b46eb9971c..fa2b6c0199 100644
--- a/packages/services/api/src/modules/schema/providers/schema-manager.ts
+++ b/packages/services/api/src/modules/schema/providers/schema-manager.ts
@@ -3,6 +3,7 @@ import { parse, print } from 'graphql';
 import { Inject, Injectable, Scope } from 'graphql-modules';
 import lodash from 'lodash';
 import { z } from 'zod';
+import { traceFn } from '@hive/service-common';
 import type {
   ConditionalBreakingChangeMetadata,
   SchemaChangeType,
@@ -95,6 +96,15 @@ export class SchemaManager {
     });
   }
 
+  @traceFn('SchemaManager.compose', {
+    initAttributes: input => ({
+      'hive.target.id': input.targetId,
+      'hive.organization.id': input.organizationId,
+      'hive.project.id': input.projectId,
+      'input.only.composable': input.onlyComposable,
+      'input.services.count': input.services.length,
+    }),
+  })
   async compose(
     input: TargetSelector & {
       onlyComposable: boolean;
@@ -187,6 +197,14 @@ export class SchemaManager {
     throw new Error('Composition was successful but is missing a supergraph');
   }
 
+  @traceFn('SchemaManager.getSchemasOfVersion', {
+    initAttributes: selector => ({
+      'hive.target.id': selector.targetId,
+      'hive.organization.id': selector.organizationId,
+      'hive.project.id': selector.projectId,
+      'hive.version.id': selector.versionId,
+    }),
+  })
   @atomic(stringifySelector)
   async getSchemasOfVersion(
     selector: {
@@ -204,6 +222,14 @@ export class SchemaManager {
     return schemas;
   }
 
+  @traceFn('SchemaManager.getSchemasOfVersion', {
+    initAttributes: input => ({
+      'hive.target.id': input.targetId,
+      'hive.organization.id': input.organizationId,
+      'hive.project.id': input.projectId,
+      'hive.version.id': input.id,
+    }),
+  })
   @atomic(stringifySelector)
   async getMaybeSchemasOfVersion(schemaVersion: SchemaVersion) {
     this.logger.debug('Fetching schemas (schemaVersionId=%s)', schemaVersion.id);
@@ -298,6 +324,15 @@ export class SchemaManager {
     };
   }
 
+  @traceFn('SchemaManager.updateSchemaVersionStatus', {
+    initAttributes: input => ({
+      'hive.target.id': input.targetId,
+      'hive.organization.id': input.organizationId,
+      'hive.project.id': input.projectId,
+      'hive.version.id': input.versionId,
+      'hive.input.valid': input.valid,
+    }),
+  })
   async updateSchemaVersionStatus(
     input: TargetSelector & { versionId: string; valid: boolean },
   ): Promise<SchemaVersion> {
@@ -337,6 +372,16 @@ export class SchemaManager {
     });
   }
 
+  @traceFn('SchemaManager.createVersion', {
+    initAttributes: input => ({
+      'hive.target.id': input.targetId,
+      'hive.organization.id': input.organizationId,
+      'hive.project.id': input.projectId,
+      'hive.version.commit': input.commit,
+      'hive.version.valid': input.valid,
+      'hive.version.service': input.service || '',
+    }),
+  })
   async createVersion(
     input: ({
       commit: string;
diff --git a/packages/services/api/src/modules/schema/providers/schema-version-helper.ts b/packages/services/api/src/modules/schema/providers/schema-version-helper.ts
index e628ab893f..a480938c74 100644
--- a/packages/services/api/src/modules/schema/providers/schema-version-helper.ts
+++ b/packages/services/api/src/modules/schema/providers/schema-version-helper.ts
@@ -2,6 +2,7 @@ import type { SchemaVersionMapper as SchemaVersion } from '../module.graphql.map
 import { print } from 'graphql';
 import { Injectable, Scope } from 'graphql-modules';
 import { CriticalityLevel } from '@graphql-inspector/core';
+import { traceFn } from '@hive/service-common';
 import type { SchemaChangeType } from '@hive/storage';
 import {
   containsSupergraphSpec,
@@ -38,6 +39,14 @@ export class SchemaVersionHelper {
     private logger: Logger,
   ) {}
 
+  @traceFn('SchemaVersionHelper.composeSchemaVersion', {
+    initAttributes: input => ({
+      'hive.target.id': input.targetId,
+      'hive.organization.id': input.organizationId,
+      'hive.project.id': input.projectId,
+      'hive.version.id': input.id,
+    }),
+  })
   @cache<SchemaVersion>(version => version.id)
   private async composeSchemaVersion(schemaVersion: SchemaVersion) {
     const [schemas, project, organization] = await Promise.all([
@@ -145,6 +154,18 @@ export class SchemaVersionHelper {
     return supergraphAst;
   }
 
+  @traceFn('SchemaVersionHelper.getSchemaChanges', {
+    initAttributes: input => ({
+      'hive.target.id': input.targetId,
+      'hive.organization.id': input.organizationId,
+      'hive.project.id': input.projectId,
+      'hive.version.id': input.id,
+    }),
+    resultAttributes: changes => ({
+      'hive.breaking-changes.count': changes?.breaking?.length,
+      'hive.safe-changes.count': changes?.safe?.length,
+    }),
+  })
   @cache<SchemaVersion>(version => version.id)
   private async getSchemaChanges(schemaVersion: SchemaVersion) {
     if (!schemaVersion.isComposable) {
@@ -282,6 +303,15 @@ export class SchemaVersionHelper {
     return !composableVersion;
   }
 
+  @traceFn('SchemaVersionHelper.getServiceSdlForPreviousVersionService', {
+    initAttributes: (schemaVersion, serviceName) => ({
+      'hive.organization.id': schemaVersion.organizationId,
+      'hive.project.id': schemaVersion.projectId,
+      'hive.target.id': schemaVersion.targetId,
+      'hive.version.id': schemaVersion.id,
+      'hive.service.name': serviceName,
+    }),
+  })
   async getServiceSdlForPreviousVersionService(schemaVersion: SchemaVersion, serviceName: string) {
     const previousVersion = await this.getPreviousDiffableSchemaVersion(schemaVersion);
     if (!previousVersion) {