Allow Distro config to be specified (#1382)

* Allow Distro config to be specified * Update test
microsoft · Sep 11, 2024 · d210010 · d210010
1 parent 2b31e1f
commit d210010
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 26 deletions.
diff --git a/src/shim/shim-config.ts b/src/shim/shim-config.ts
@@ -61,6 +61,9 @@ class Config implements IConfig {
     public noPatchModules: string;
     public noDiagnosticChannel: boolean;
 
+    // Expose Distro config for further customization, other conflicting configs will take precedence over this.
+    public azureMonitorOpenTelemetryOptions : AzureMonitorOpenTelemetryOptions;
+
     private _configWarnings: string[];
 
     /**
@@ -124,29 +127,33 @@ class Config implements IConfig {
     * Parse the config property to set the appropriate values on the AzureMonitorOpenTelemetryOptions
     */
     public parseConfig(): AzureMonitorOpenTelemetryOptions {
-        const options: AzureMonitorOpenTelemetryOptions = {
-            azureMonitorExporterOptions: {
-                connectionString: this.connectionString,
-                disableOfflineStorage: false,
-            },
-            enableAutoCollectPerformance: true,
-            enableAutoCollectExceptions: true,
-            instrumentationOptions: {
-                http: { enabled: true },
-                azureSdk: { enabled: true },
-                mongoDb: { enabled: true },
-                mySql: { enabled: true },
-                redis: { enabled: true },
-                redis4: { enabled: true },
-                postgreSql: { enabled: true },
-                bunyan: { enabled: true },
-                winston: { enabled: true },
+        const options: AzureMonitorOpenTelemetryOptions = Object.assign(
+            {
+                azureMonitorExporterOptions: {
+                    connectionString: this.connectionString,
+                    disableOfflineStorage: false,
+                },
+                enableAutoCollectPerformance: true,
+                enableAutoCollectExceptions: true,
+                instrumentationOptions: {
+                    http: { enabled: true },
+                    azureSdk: { enabled: true },
+                    mongoDb: { enabled: true },
+                    mySql: { enabled: true },
+                    redis: { enabled: true },
+                    redis4: { enabled: true },
+                    postgreSql: { enabled: true },
+                    bunyan: { enabled: true },
+                    winston: { enabled: true },
+                },
+                otlpTraceExporterConfig: {},
+                otlpMetricExporterConfig: {},
+                otlpLogExporterConfig: {},
+                enableLiveMetrics: true,
             },
-            otlpTraceExporterConfig: {},
-            otlpMetricExporterConfig: {},
-            otlpLogExporterConfig: {},
-            enableLiveMetrics: true,
-        };
+            this.azureMonitorOpenTelemetryOptions
+        );
+
         (options.instrumentationOptions as InstrumentationOptions) = {
             ...options.instrumentationOptions,
             console: { enabled: false },

diff --git a/test/unitTests/shim/config.tests.ts b/test/unitTests/shim/config.tests.ts
@@ -10,6 +10,9 @@ import http = require("http");
 import https = require("https");
 import { DistributedTracingModes } from '../../../applicationinsights';
 import { checkWarnings } from './testUtils';
+import { BatchLogRecordProcessor, ConsoleLogRecordExporter } from '@opentelemetry/sdk-logs';
+import { BatchSpanProcessor, ConsoleSpanExporter } from '@opentelemetry/sdk-trace-base';
+import { Resource } from '@opentelemetry/resources';
 
 class TestTokenCredential implements azureCoreAuth.TokenCredential {
     private _expiresOn: Date;
@@ -81,7 +84,7 @@ describe("shim/configuration/config", () => {
                 "winston": { "enabled": true },
                 "console": { "enabled": true },
             }),
-            "wrong instrumentationOptions");
+                "wrong instrumentationOptions");
             assert.equal(JSON.stringify(options.instrumentationOptions.bunyan), JSON.stringify({ enabled: true }), "wrong bunyan setting");
             assert.equal(options.enableAutoCollectExceptions, true, "wrong enableAutoCollectExceptions");
             assert.equal(options.enableAutoCollectPerformance, true, "wrong enableAutoCollectPerformance");
@@ -111,6 +114,31 @@ describe("shim/configuration/config", () => {
             assert.equal(options.samplingRatio, 0, "wrong samplingRatio");
         });
 
+
+        it("should allow customization of Azure Monitor Distro configuration", () => {
+            let spanProcessors = [new BatchSpanProcessor(new ConsoleSpanExporter())];
+            let logRecordProcessors = [new BatchLogRecordProcessor(new ConsoleLogRecordExporter)];
+            let resource = new Resource({});
+            const config = new Config(connectionString);
+            config.azureMonitorOpenTelemetryOptions = {
+                resource: resource,
+                enableTraceBasedSamplingForLogs: false,
+                enableLiveMetrics: false,
+                enableStandardMetrics: false,
+                logRecordProcessors: logRecordProcessors,
+                spanProcessors: spanProcessors
+            };
+
+            let options = config.parseConfig();
+            assert.equal(options.resource, resource, "wrong resource");
+            assert.equal(options.enableTraceBasedSamplingForLogs, false, "wrong enableTraceBasedSamplingForLogs");
+            assert.equal(options.enableLiveMetrics, false, "wrong enableTraceBasedSamplingForLogs");
+            assert.equal(options.enableStandardMetrics, false, "wrong enableTraceBasedSamplingForLogs");
+            assert.equal(options.logRecordProcessors, logRecordProcessors, "wrong logRecordProcessors");
+            assert.equal(options.spanProcessors, spanProcessors, "wrong spanProcessors");
+        });
+
+
         it("should activate DEBUG internal logger", () => {
             const env = <{ [id: string]: string }>{};
             process.env = env;
@@ -151,15 +179,15 @@ describe("shim/configuration/config", () => {
                 "postgreSql": { "enabled": false },
                 "bunyan": { "enabled": false },
                 "winston": { "enabled": false },
-                "console":{ "enabled": false },
+                "console": { "enabled": false },
             }));
         });
 
         it("should disable specific instrumentations when noPatchModules is set", () => {
             const config = new Config(connectionString);
             config.noPatchModules = "azuresdk,mongodb-core,redis,pg-pool";
             let options = config.parseConfig();
-            assert.equal(JSON.stringify(options.instrumentationOptions), JSON.stringify({ 
+            assert.equal(JSON.stringify(options.instrumentationOptions), JSON.stringify({
                 http: { enabled: true },
                 azureSdk: { enabled: false },
                 mongoDb: { enabled: false },
@@ -326,7 +354,7 @@ describe("shim/configuration/config", () => {
             it("should warn if web instrumentations are set", () => {
                 const config = new Config(connectionString);
                 const warnings = config["_configWarnings"];
-                config.webInstrumentationConfig = [{name: "test", value: true}];
+                config.webInstrumentationConfig = [{ name: "test", value: true }];
                 config.webInstrumentationSrc = "test";
                 config.parseConfig();
                 assert.ok(checkWarnings("The webInstrumentation config and src options are not supported by the shim.", warnings), "warning was not raised");