From 587b032d103f94916cbce67fd699e5222df5a8a7 Mon Sep 17 00:00:00 2001
From: Patryk Ziemkowski <pziemkowski@apptension.com>
Date: Thu, 14 Sep 2023 12:23:10 +0200
Subject: [PATCH] Modify base command to support nested telemetry spans

---
 packages/internal/cli/src/baseCommand.ts      | 57 ++++++++++++++-----
 packages/internal/cli/src/config/telemetry.ts |  4 ++
 2 files changed, 47 insertions(+), 14 deletions(-)

diff --git a/packages/internal/cli/src/baseCommand.ts b/packages/internal/cli/src/baseCommand.ts
index 001334a0b..ab3435812 100644
--- a/packages/internal/cli/src/baseCommand.ts
+++ b/packages/internal/cli/src/baseCommand.ts
@@ -2,8 +2,7 @@ import { Command, Flags, Interfaces } from '@oclif/core';
 import { ExitError } from '@oclif/core/lib/errors';
 import { Span, SpanStatusCode, trace, Tracer } from '@opentelemetry/api';
 
-import { SB_TELEMETRY_DISABLED } from './config/env';
-import { traceExporter } from './config/telemetry';
+import * as telemetry from './config/telemetry';
 
 const formatAttrs = (obj: { [k: string]: string } = {}, prefix = '') => {
   return Object.fromEntries(
@@ -38,21 +37,52 @@ export abstract class BaseCommand<T extends typeof Command> extends Command {
     this.flags = flags as Flags<T>;
     this.args = args as Args<T>;
 
-    if (!SB_TELEMETRY_DISABLED) {
+    if (telemetry.isEnabled) {
       this.printTelemetryInfo();
       this.tracer = trace.getTracer('command', this.config.version);
+    }
+  }
 
-      this.span = this.tracer.startSpan(`command.${this.ctor.id}`, {
-        attributes: {
-          ...formatAttrs(flags, 'flags'),
-          ...formatAttrs(args, 'args'),
-        },
-      });
+  async _run<T>() {
+    let err;
+    let result;
+    try {
+      // remove redirected env var to allow subsessions to run autoupdated client
+      // @ts-ignore
+      this.removeEnvVar('REDIRECTED');
+      await this.init();
+
+      if (telemetry.isEnabled && this.tracer) {
+        result = this.tracer.startActiveSpan(
+          `command.${this.ctor.id}`,
+          {
+            attributes: {
+              ...formatAttrs(this.flags, 'flags'),
+              ...formatAttrs(this.args, 'args'),
+            },
+          },
+          async (span) => {
+            this.span = span;
+            const _result = this.run();
+            span.end();
+            return _result;
+          }
+        );
+      } else {
+        result = await this.run();
+      }
+    } catch (error: any) {
+      err = error;
+      await this.catch(error);
+    } finally {
+      await this.finally(err);
     }
+    if (result && this.jsonEnabled()) this.logJson(this.toSuccessJson(result));
+    return result;
   }
 
   protected async catch(err: Error & { exitCode?: number }): Promise<any> {
-    if (!SB_TELEMETRY_DISABLED) {
+    if (telemetry.isEnabled) {
       if (!(err instanceof ExitError) || err.oclif.exit !== 0) {
         this.span?.addEvent('Command error');
         this.span?.recordException(err);
@@ -63,21 +93,20 @@ export abstract class BaseCommand<T extends typeof Command> extends Command {
   }
 
   protected async finally(_: Error | undefined): Promise<any> {
-    if (!SB_TELEMETRY_DISABLED) {
+    if (telemetry.isEnabled) {
       this.span?.addEvent('Command finished');
       this.span?.end();
 
       // Need to wait en event loop for the internal promise in exporter to be visible
       await new Promise((resolve) => setTimeout(() => resolve(true)));
       // wait for the exporter to send data
-      await traceExporter.forceFlush();
+      await telemetry.traceExporter.forceFlush();
     }
     return super.finally(_);
   }
 
   protected printTelemetryInfo(): void {
-    console.log({SB_TELEMETRY_DISABLED})
-    if (!SB_TELEMETRY_DISABLED) {
+    if (telemetry.isEnabled) {
       this.log(`\x1b[2m
 ------ Notice ------
 This CLI collects various anonymous events, warnings, and errors to improve the CLI tool and enhance your user experience.
diff --git a/packages/internal/cli/src/config/telemetry.ts b/packages/internal/cli/src/config/telemetry.ts
index 4a029b2ae..a14babbe7 100644
--- a/packages/internal/cli/src/config/telemetry.ts
+++ b/packages/internal/cli/src/config/telemetry.ts
@@ -4,8 +4,12 @@ import { NodeTracerProvider } from '@opentelemetry/sdk-trace-node';
 import { SemanticResourceAttributes } from '@opentelemetry/semantic-conventions';
 import { OTLPTraceExporter } from '@opentelemetry/exporter-trace-otlp-proto';
 
+import { SB_TELEMETRY_DISABLED } from './env';
+
 const sbTelemetry = require('@apptension/saas-boilerplate-telemetry');
 
+export const isEnabled = !SB_TELEMETRY_DISABLED;
+
 export const traceExporter = new OTLPTraceExporter({
   url: sbTelemetry?.[0],
   headers: {