feat: telemetry (#216)

common/lib/abstract_connection_plugin.ts

@@ -23,6 +23,7 @@ import { HostRole } from "./host_role";
 import { ClientWrapper } from "./client_wrapper";
 
 export abstract class AbstractConnectionPlugin implements ConnectionPlugin {
+  name: string = this.constructor.name;
   abstract getSubscribedMethods(): Set<string>;
 
   connect(

common/lib/authentication/aws_secrets_manager_plugin.ts

@@ -15,10 +15,10 @@
 */
 
 import {
-  SecretsManagerClientConfig,
+  GetSecretValueCommand,
   SecretsManagerClient,
-  SecretsManagerServiceException,
-  GetSecretValueCommand
+  SecretsManagerClientConfig,
+  SecretsManagerServiceException
 } from "@aws-sdk/client-secrets-manager";
 import { logger } from "../../logutils";
 import { AbstractConnectionPlugin } from "../abstract_connection_plugin";
@@ -29,11 +29,15 @@ import { Messages } from "../utils/messages";
 import { WrapperProperties } from "../wrapper_property";
 import { logAndThrowError } from "../utils/utils";
 import { ClientWrapper } from "../client_wrapper";
+import { TelemetryTraceLevel } from "../utils/telemetry/telemetry_trace_level";
 
 export class AwsSecretsManagerPlugin extends AbstractConnectionPlugin {
+  private static readonly TELEMETRY_UPDATE_SECRETS = "fetch credentials";
+  private static readonly TELEMETRY_FETCH_CREDENTIALS_COUNTER = "secretsManager.fetchCredentials.count";
   private static SUBSCRIBED_METHODS: Set<string> = new Set<string>(["connect", "forceConnect"]);
   private static SECRETS_ARN_PATTERN: RegExp = new RegExp("^arn:aws:secretsmanager:(?<region>[^:\\n]*):[^:\\n]*:([^:/\\n]*[:/])?(.*)$");
   private readonly pluginService: PluginService;
+  private readonly fetchCredentialsCounter;
   private secret: Secret | null = null;
   static secretsCache: Map<string, Secret> = new Map();
   secretKey: SecretCacheKey;
@@ -66,6 +70,10 @@ export class AwsSecretsManagerPlugin extends AbstractConnectionPlugin {
 
     this.secretKey = new SecretCacheKey(secretId, region);
     this.secretsManagerClient = new SecretsManagerClient(config);
+
+    this.fetchCredentialsCounter = this.pluginService
+      .getTelemetryFactory()
+      .createCounter(AwsSecretsManagerPlugin.TELEMETRY_FETCH_CREDENTIALS_COUNTER);
   }
 
   getSubscribedMethods(): Set<string> {
@@ -117,26 +125,31 @@ export class AwsSecretsManagerPlugin extends AbstractConnectionPlugin {
   }
 
   private async updateSecret(forceRefresh: boolean): Promise<boolean> {
-    let fetched = false;
-    this.secret = AwsSecretsManagerPlugin.secretsCache.get(JSON.stringify(this.secretKey)) ?? null;
-
-    if (!this.secret || forceRefresh) {
-      try {
-        this.secret = await this.fetchLatestCredentials();
-        fetched = true;
-        AwsSecretsManagerPlugin.secretsCache.set(JSON.stringify(this.secretKey), this.secret);
-      } catch (error: any) {
-        if (error instanceof SecretsManagerServiceException) {
-          logAndThrowError(Messages.get("AwsSecretsManagerConnectionPlugin.failedToFetchDbCredentials"));
-        } else if (error instanceof Error && error.message.includes("AWS SDK error")) {
-          logAndThrowError(Messages.get("AwsSecretsManagerConnectionPlugin.endpointOverrideInvalidConnection", error.message));
-        } else {
-          logAndThrowError(Messages.get("AwsSecretsManagerConnectionPlugin.unhandledException", error.message));
+    const telemetryFactory = this.pluginService.getTelemetryFactory();
+    const telemetryContext = telemetryFactory.openTelemetryContext(AwsSecretsManagerPlugin.TELEMETRY_UPDATE_SECRETS, TelemetryTraceLevel.NESTED);
+    this.fetchCredentialsCounter.inc();
+
+    return await telemetryContext.start(async () => {
+      let fetched = false;
+      this.secret = AwsSecretsManagerPlugin.secretsCache.get(JSON.stringify(this.secretKey)) ?? null;
+
+      if (!this.secret || forceRefresh) {
+        try {
+          this.secret = await this.fetchLatestCredentials();
+          fetched = true;
+          AwsSecretsManagerPlugin.secretsCache.set(JSON.stringify(this.secretKey), this.secret);
+        } catch (error: any) {
+          if (error instanceof SecretsManagerServiceException) {
+            logAndThrowError(Messages.get("AwsSecretsManagerConnectionPlugin.failedToFetchDbCredentials"));
+          } else if (error instanceof Error && error.message.includes("AWS SDK error")) {
+            logAndThrowError(Messages.get("AwsSecretsManagerConnectionPlugin.endpointOverrideInvalidConnection", error.message));
+          } else {
+            logAndThrowError(Messages.get("AwsSecretsManagerConnectionPlugin.unhandledException", error.message));
+          }
         }
       }
-    }
-
-    return fetched;
+      return fetched;
+    });
   }
 
   private async fetchLatestCredentials(): Promise<Secret> {

common/lib/authentication/iam_authentication_plugin.ts

@@ -25,16 +25,22 @@ import { AbstractConnectionPlugin } from "../abstract_connection_plugin";
 import { WrapperProperties } from "../wrapper_property";
 import { IamAuthUtils, TokenInfo } from "../utils/iam_auth_utils";
 import { ClientWrapper } from "../client_wrapper";
+import { AwsCredentialIdentityProvider } from "@smithy/types/dist-types/identity/awsCredentialIdentity";
+import { TelemetryTraceLevel } from "../utils/telemetry/telemetry_trace_level";
 
 export class IamAuthenticationPlugin extends AbstractConnectionPlugin {
   private static readonly SUBSCRIBED_METHODS = new Set<string>(["connect", "forceConnect"]);
   protected static readonly tokenCache = new Map<string, TokenInfo>();
+  private readonly telemetryFactory;
+  private readonly fetchTokenCounter;
   private pluginService: PluginService;
   rdsUtil: RdsUtils = new RdsUtils();
 
   constructor(pluginService: PluginService) {
     super();
     this.pluginService = pluginService;
+    this.telemetryFactory = this.pluginService.getTelemetryFactory();
+    this.fetchTokenCounter = this.telemetryFactory.createCounter("iam.fetchTokenCount");
   }
 
   getSubscribedMethods(): Set<string> {
@@ -89,9 +95,11 @@ export class IamAuthenticationPlugin extends AbstractConnectionPlugin {
         host,
         port,
         region,
-        WrapperProperties.USER.get(props),
-        AwsCredentialsManager.getProvider(hostInfo, props)
+        user,
+        AwsCredentialsManager.getProvider(hostInfo, props),
+        this.pluginService
       );
+      this.fetchTokenCounter.inc();
       logger.debug(Messages.get("AuthenticationToken.generatedNewToken", token));
       WrapperProperties.PASSWORD.set(props, token);
       IamAuthenticationPlugin.tokenCache.set(cacheKey, new TokenInfo(token, tokenExpiry));
@@ -114,9 +122,11 @@ export class IamAuthenticationPlugin extends AbstractConnectionPlugin {
         host,
         port,
         region,
-        WrapperProperties.USER.get(props),
-        AwsCredentialsManager.getProvider(hostInfo, props)
+        user,
+        AwsCredentialsManager.getProvider(hostInfo, props),
+        this.pluginService
       );
+      this.fetchTokenCounter.inc();
       logger.debug(Messages.get("AuthenticationToken.generatedNewToken", token));
       WrapperProperties.PASSWORD.set(props, token);
       IamAuthenticationPlugin.tokenCache.set(cacheKey, new TokenInfo(token, tokenExpiry));

common/lib/aws_client.ts

@@ -24,9 +24,12 @@ import { PluginManager } from "./plugin_manager";
 import { EventEmitter } from "stream";
 import { DriverConnectionProvider } from "./driver_connection_provider";
 import { ClientWrapper } from "./client_wrapper";
+import { DefaultTelemetryFactory } from "./utils/telemetry/default_telemetry_factory";
+import { TelemetryFactory } from "./utils/telemetry/telemetry_factory";
 
 export abstract class AwsClient extends EventEmitter {
   private _defaultPort: number = -1;
+  protected telemetryFactory: TelemetryFactory;
   protected pluginManager: PluginManager;
   protected pluginService: PluginService;
   protected isConnected: boolean = false;
@@ -60,12 +63,15 @@ export abstract class AwsClient extends EventEmitter {
     const effectiveConnProvider = null;
     // TODO: check for configuration profile to update the effectiveConnProvider
 
+    this.telemetryFactory = new DefaultTelemetryFactory(this.properties);
+
     const container = new PluginServiceManagerContainer();
     this.pluginService = new PluginService(container, this, dbType, knownDialectsByCode, this.properties);
-    this.pluginManager = new PluginManager(container, this.properties, defaultConnProvider, effectiveConnProvider);
+    this.pluginManager = new PluginManager(container, this.properties, defaultConnProvider, effectiveConnProvider, this.telemetryFactory);
   }
 
   private async setup() {
+    await this.telemetryFactory.init();
     await this.pluginManager.init();
   }
 
@@ -88,6 +94,7 @@ export abstract class AwsClient extends EventEmitter {
     if (info != null) {
       await this.pluginService.refreshHostList();
     }
+
     this.isConnected = true;
   }
 

common/lib/connection_plugin.ts

@@ -22,6 +22,7 @@ import { HostRole } from "./host_role";
 import { ClientWrapper } from "./client_wrapper";
 
 export interface ConnectionPlugin {
+  name: string;
   getSubscribedMethods(): Set<string>;
 
   connect(

common/lib/connection_provider.ts

@@ -25,4 +25,5 @@ export interface ConnectionProvider {
   acceptsUrl(hostInfo: HostInfo, props: Map<string, any>): boolean;
   acceptsStrategy(role: HostRole, strategy: string): boolean;
   getHostInfoByStrategy(hosts: HostInfo[], role: HostRole, strategy: string, props?: Map<string, any>): HostInfo;
+  getTargetName(): string;
 }

common/lib/driver_connection_provider.ts

@@ -134,4 +134,8 @@ export class DriverConnectionProvider implements ConnectionProvider {
     }
     return acceptedStrategy.getHost(hosts, role, props);
   }
+
+  getTargetName(): string {
+    return this.constructor.name;
+  }
 }

common/lib/host_list_provider_service.ts

@@ -20,6 +20,7 @@ import { AwsClient } from "./aws_client";
 import { DatabaseDialect } from "./database_dialect/database_dialect";
 import { HostInfoBuilder } from "./host_info_builder";
 import { ConnectionUrlParser } from "./utils/connection_url_parser";
+import { TelemetryFactory } from "./utils/telemetry/telemetry_factory";
 
 export interface HostListProviderService {
   getHostListProvider(): HostListProvider | null;
@@ -47,4 +48,8 @@ export interface HostListProviderService {
   setInTransaction(inTransaction: boolean): void;
 
   isClientValid(targetClient: any): Promise<boolean>;
+
+  getTelemetryFactory(): TelemetryFactory;
+
+  getTargetName(): string;
 }