telemetry (#414)

Co-authored-by: justing-bq <[email protected]>
Co-authored-by: sergiyv-bitquill <[email protected]>

Author: 3 people

aws_advanced_python_wrapper/aws_secrets_manager_plugin.py

@@ -35,6 +35,8 @@
 from aws_advanced_python_wrapper.utils.messages import Messages
 from aws_advanced_python_wrapper.utils.properties import (Properties,
                                                           WrapperProperties)
+from aws_advanced_python_wrapper.utils.telemetry.telemetry import \
+    TelemetryTraceLevel
 
 logger = Logger(__name__)
 
@@ -66,6 +68,9 @@ def __init__(self, plugin_service: PluginService, props: Properties, session: Op
         secrets_endpoint = WrapperProperties.SECRETS_MANAGER_ENDPOINT.get(props)
         self._secret_key: Tuple = (secret_id, region, secrets_endpoint)
 
+        telemetry_factory = self._plugin_service.get_telemetry_factory()
+        self._fetch_credentials_counter = telemetry_factory.create_counter("secrets_manager.fetch_credentials.count")
+
     def connect(
             self,
             target_driver_func: Callable,
@@ -116,35 +121,44 @@ def _update_secret(self, force_refetch: bool = False) -> bool:
         :param force_refetch: Allows ignoring cached credentials and force fetches the latest credentials from the service.
         :return: `True`, if credentials were fetched from the service.
         """
-        fetched: bool = False
-
-        self._secret: Optional[SimpleNamespace] = AwsSecretsManagerPlugin._secrets_cache.get(self._secret_key)
-        endpoint = self._secret_key[2]
-
-        if not self._secret or force_refetch:
-            try:
-                self._secret = self._fetch_latest_credentials()
-                if self._secret:
-                    AwsSecretsManagerPlugin._secrets_cache[self._secret_key] = self._secret
-                    fetched = True
-            except (ClientError, AttributeError) as e:
-                logger.debug("AwsSecretsManagerPlugin.FailedToFetchDbCredentials", e)
-                raise AwsWrapperError(
-                    Messages.get_formatted("AwsSecretsManagerPlugin.FailedToFetchDbCredentials", e)) from e
-            except JSONDecodeError as e:
-                logger.debug("AwsSecretsManagerPlugin.JsonDecodeError", e)
-                raise AwsWrapperError(
-                    Messages.get_formatted("AwsSecretsManagerPlugin.JsonDecodeError", e))
-            except EndpointConnectionError:
-                logger.debug("AwsSecretsManagerPlugin.EndpointOverrideInvalidConnection", endpoint)
-                raise AwsWrapperError(
-                    Messages.get_formatted("AwsSecretsManagerPlugin.EndpointOverrideInvalidConnection", endpoint))
-            except ValueError:
-                logger.debug("AwsSecretsManagerPlugin.EndpointOverrideMisconfigured", endpoint)
-                raise AwsWrapperError(
-                    Messages.get_formatted("AwsSecretsManagerPlugin.EndpointOverrideMisconfigured", endpoint))
+        telemetry_factory = self._plugin_service.get_telemetry_factory()
+        context = telemetry_factory.open_telemetry_context("fetch credentials", TelemetryTraceLevel.NESTED)
+        self._fetch_credentials_counter.inc()
 
-        return fetched
+        try:
+            fetched: bool = False
+            self._secret: Optional[SimpleNamespace] = AwsSecretsManagerPlugin._secrets_cache.get(self._secret_key)
+            endpoint = self._secret_key[2]
+            if not self._secret or force_refetch:
+                try:
+                    self._secret = self._fetch_latest_credentials()
+                    if self._secret:
+                        AwsSecretsManagerPlugin._secrets_cache[self._secret_key] = self._secret
+                        fetched = True
+                except (ClientError, AttributeError) as e:
+                    logger.debug("AwsSecretsManagerPlugin.FailedToFetchDbCredentials", e)
+                    raise AwsWrapperError(
+                        Messages.get_formatted("AwsSecretsManagerPlugin.FailedToFetchDbCredentials", e)) from e
+                except JSONDecodeError as e:
+                    logger.debug("AwsSecretsManagerPlugin.JsonDecodeError", e)
+                    raise AwsWrapperError(
+                        Messages.get_formatted("AwsSecretsManagerPlugin.JsonDecodeError", e))
+                except EndpointConnectionError:
+                    logger.debug("AwsSecretsManagerPlugin.EndpointOverrideInvalidConnection", endpoint)
+                    raise AwsWrapperError(
+                        Messages.get_formatted("AwsSecretsManagerPlugin.EndpointOverrideInvalidConnection", endpoint))
+                except ValueError:
+                    logger.debug("AwsSecretsManagerPlugin.EndpointOverrideMisconfigured", endpoint)
+                    raise AwsWrapperError(
+                        Messages.get_formatted("AwsSecretsManagerPlugin.EndpointOverrideMisconfigured", endpoint))
+
+            return fetched
+        except Exception as ex:
+            context.set_success(False)
+            context.set_exception(ex)
+            raise ex
+        finally:
+            context.close_context()
 
     def _fetch_latest_credentials(self):
         """

aws_advanced_python_wrapper/default_plugin.py

@@ -33,6 +33,8 @@
 from aws_advanced_python_wrapper.hostinfo import HostInfo, HostRole
 from aws_advanced_python_wrapper.plugin import Plugin
 from aws_advanced_python_wrapper.utils.messages import Messages
+from aws_advanced_python_wrapper.utils.telemetry.telemetry import \
+    TelemetryTraceLevel
 
 
 class DefaultPlugin(Plugin):
@@ -64,11 +66,20 @@ def _connect(
             host_info: HostInfo,
             props: Properties,
             conn_provider: ConnectionProvider) -> Connection:
-        database_dialect = self._plugin_service.database_dialect
-        conn = conn_provider.connect(target_func, driver_dialect, database_dialect, host_info, props)
+        telemetry_factory = self._plugin_service.get_telemetry_factory()
+        context = telemetry_factory.open_telemetry_context(driver_dialect.driver_name, TelemetryTraceLevel.NESTED)
+
+        conn: Connection
+        try:
+            database_dialect = self._plugin_service.database_dialect
+            conn = conn_provider.connect(target_func, driver_dialect, database_dialect, host_info, props)
+        finally:
+            context.close_context()
+
         self._plugin_service.set_availability(host_info.all_aliases, HostAvailability.AVAILABLE)
         self._plugin_service.update_driver_dialect(conn_provider)
         self._plugin_service.update_dialect(conn)
+
         return conn
 
     def force_connect(
@@ -88,7 +99,15 @@ def force_connect(
             self._connection_provider_manager.default_provider)
 
     def execute(self, target: object, method_name: str, execute_func: Callable, *args: Any, **kwargs: Any) -> Any:
-        result = self._plugin_service.driver_dialect.execute(method_name, execute_func, *args, **kwargs)
+        telemetry_factory = self._plugin_service.get_telemetry_factory()
+        context = telemetry_factory.open_telemetry_context(
+            self._plugin_service.driver_dialect.driver_name, TelemetryTraceLevel.NESTED)
+
+        try:
+            result = self._plugin_service.driver_dialect.execute(method_name, execute_func, *args, **kwargs)
+        finally:
+            context.close_context()
+
         if method_name != DefaultPlugin._CLOSE_METHOD and self._plugin_service.current_connection is not None:
             self._plugin_service.update_in_transaction()
 

aws_advanced_python_wrapper/failover_plugin.py

@@ -43,6 +43,8 @@
                                                           WrapperProperties)
 from aws_advanced_python_wrapper.utils.rds_url_type import RdsUrlType
 from aws_advanced_python_wrapper.utils.rdsutils import RdsUtils
+from aws_advanced_python_wrapper.utils.telemetry.telemetry import \
+    TelemetryTraceLevel
 from aws_advanced_python_wrapper.writer_failover_handler import (
     WriterFailoverHandler, WriterFailoverHandlerImpl)
 
@@ -84,7 +86,10 @@ def __init__(self, plugin_service: PluginService, props: Properties):
             self._properties)
         self._failover_reader_connect_timeout_sec = WrapperProperties.FAILOVER_READER_CONNECT_TIMEOUT_SEC.get_float(
             self._properties)
-        self._keep_session_state_on_failover = WrapperProperties.KEEP_SESSION_STATE_ON_FAILOVER.get_bool(self._properties)
+        self._keep_session_state_on_failover = WrapperProperties.KEEP_SESSION_STATE_ON_FAILOVER.get_bool(
+            self._properties)
+        self._telemetry_failover_additional_top_trace_setting = (
+            WrapperProperties.TELEMETRY_FAILOVER_ADDITIONAL_TOP_TRACE.get_bool(self._properties))
         self._failover_mode: FailoverMode
         self._is_in_transaction: bool = False
         self._is_closed: bool = False
@@ -96,6 +101,18 @@ def __init__(self, plugin_service: PluginService, props: Properties):
         self._saved_read_only_status: bool = False
         self._saved_auto_commit_status: bool = False
 
+        telemetry_factory = self._plugin_service.get_telemetry_factory()
+        self._failover_writer_triggered_counter = telemetry_factory.create_counter("writer_failover.triggered.count")
+        self._failover_writer_success_counter = telemetry_factory.create_counter(
+            "writer_failover.completed.success.count")
+        self._failover_writer_failed_counter = telemetry_factory.create_counter(
+            "writer_failover.completed.failed.count")
+        self._failover_reader_triggered_counter = telemetry_factory.create_counter("reader_failover.triggered.count")
+        self._failover_reader_success_counter = telemetry_factory.create_counter(
+            "reader_failover.completed.success.count")
+        self._failover_reader_failed_counter = telemetry_factory.create_counter(
+            "reader_failover.completed.failed.count")
+
         FailoverPlugin._SUBSCRIBED_METHODS.update(self._plugin_service.network_bound_methods)
 
     def init_host_provider(
@@ -213,10 +230,13 @@ def _connect(
             properties: Properties,
             is_initial_connection: bool,
             connect_func: Callable) -> Connection:
-        conn: Connection = self._stale_dns_helper.get_verified_connection(is_initial_connection, self._host_list_provider_service, host, properties,
+        conn: Connection = self._stale_dns_helper.get_verified_connection(is_initial_connection,
+                                                                          self._host_list_provider_service, host,
+                                                                          properties,
                                                                           connect_func)
         if self._keep_session_state_on_failover:
-            self._saved_read_only_status = False if self._saved_read_only_status == self._plugin_service.driver_dialect.is_read_only(conn) \
+            self._saved_read_only_status = False if self._saved_read_only_status == self._plugin_service.driver_dialect.is_read_only(
+                conn) \
                 else self._saved_read_only_status
             self._saved_auto_commit_status = False \
                 if self._saved_read_only_status == self._plugin_service.driver_dialect.get_autocommit(conn) \
@@ -270,53 +290,96 @@ def _failover(self, failed_host: Optional[HostInfo]):
             raise FailoverSuccessError(Messages.get(error_msg))
 
     def _failover_reader(self, failed_host: Optional[HostInfo]):
-        logger.debug("FailoverPlugin.StartReaderFailover")
+        telemetry_factory = self._plugin_service.get_telemetry_factory()
+        context = telemetry_factory.open_telemetry_context("failover to replica", TelemetryTraceLevel.NESTED)
+        self._failover_reader_triggered_counter.inc()
 
-        old_aliases = None
-        if self._plugin_service.current_host_info is not None:
-            old_aliases = self._plugin_service.current_host_info.aliases
+        try:
+            logger.debug("FailoverPlugin.StartReaderFailover")
 
-        if failed_host is not None and failed_host.get_raw_availability() != HostAvailability.AVAILABLE:
-            failed_host = None
+            old_aliases = None
+            if self._plugin_service.current_host_info is not None:
+                old_aliases = self._plugin_service.current_host_info.aliases
 
-        result: ReaderFailoverResult = self._reader_failover_handler.failover(self._plugin_service.hosts, failed_host)
+            if failed_host is not None and failed_host.get_raw_availability() != HostAvailability.AVAILABLE:
+                failed_host = None
 
-        if result is None or not result.is_connected:
-            raise FailoverFailedError(Messages.get("FailoverPlugin.UnableToConnectToReader"))
-        else:
-            if result.exception is not None:
-                raise result.exception
-            if self._keep_session_state_on_failover:
-                self.restore_session_state(result.connection)
-            if result.connection is not None and result.new_host is not None:
-                self._plugin_service.set_current_connection(result.connection, result.new_host)
+            result: ReaderFailoverResult = self._reader_failover_handler.failover(self._plugin_service.hosts,
+                                                                                  failed_host)
 
-        if self._plugin_service.current_host_info is not None and old_aliases is not None and len(old_aliases) > 0:
-            self._plugin_service.current_host_info.remove_alias(old_aliases)
+            if result is None or not result.is_connected:
+                raise FailoverFailedError(Messages.get("FailoverPlugin.UnableToConnectToReader"))
+            else:
+                if result.exception is not None:
+                    raise result.exception
+                if self._keep_session_state_on_failover:
+                    self.restore_session_state(result.connection)
+                if result.connection is not None and result.new_host is not None:
+                    self._plugin_service.set_current_connection(result.connection, result.new_host)
+
+            if self._plugin_service.current_host_info is not None and old_aliases is not None and len(old_aliases) > 0:
+                self._plugin_service.current_host_info.remove_alias(old_aliases)
+
+            self._update_topology(True)
+
+            logger.debug("FailoverPlugin.EstablishedConnection", self._plugin_service.current_host_info)
+
+            self._failover_reader_success_counter.inc()
+        except FailoverSuccessError as fse:
+            context.set_success(True)
+            context.set_exception(fse)
+            self._failover_reader_success_counter.inc()
+            raise fse
+        except Exception as ex:
+            context.set_success(False)
+            context.set_exception(ex)
+            self._failover_reader_failed_counter.inc()
+            raise ex
+        finally:
+            context.close_context()
+            if self._telemetry_failover_additional_top_trace_setting:
+                telemetry_factory.post_copy(context, TelemetryTraceLevel.FORCE_TOP_LEVEL)
 
-        self._update_topology(True)
+    def _failover_writer(self):
+        telemetry_factory = self._plugin_service.get_telemetry_factory()
+        context = telemetry_factory.open_telemetry_context("failover to writer node", TelemetryTraceLevel.NESTED)
+        self._failover_writer_triggered_counter.inc()
 
-        logger.debug("FailoverPlugin.EstablishedConnection", self._plugin_service.current_host_info)
+        try:
+            logger.debug("FailoverPlugin.StartWriterFailover")
 
-    def _failover_writer(self):
-        logger.debug("FailoverPlugin.StartWriterFailover")
+            result: WriterFailoverResult = self._writer_failover_handler.failover(self._plugin_service.hosts)
 
-        result: WriterFailoverResult = self._writer_failover_handler.failover(self._plugin_service.hosts)
+            if result is not None and result.exception is not None:
+                raise result.exception
+            elif result is None or not result.is_connected:
+                raise FailoverFailedError(Messages.get("FailoverPlugin.UnableToConnectToWriter"))
 
-        if result is not None and result.exception is not None:
-            raise result.exception
-        elif result is None or not result.is_connected:
-            raise FailoverFailedError(Messages.get("FailoverPlugin.UnableToConnectToWriter"))
+            writer_host = self._get_writer(result.topology)
+            if self._keep_session_state_on_failover:
+                self.restore_session_state(result.new_connection)
 
-        writer_host = self._get_writer(result.topology)
-        if self._keep_session_state_on_failover:
-            self.restore_session_state(result.new_connection)
+            self._plugin_service.set_current_connection(result.new_connection, writer_host)
 
-        self._plugin_service.set_current_connection(result.new_connection, writer_host)
+            logger.debug("FailoverPlugin.EstablishedConnection", self._plugin_service.current_host_info)
 
-        logger.debug("FailoverPlugin.EstablishedConnection", self._plugin_service.current_host_info)
+            self._plugin_service.refresh_host_list()
 
-        self._plugin_service.refresh_host_list()
+            self._failover_writer_success_counter.inc()
+        except FailoverSuccessError as fse:
+            context.set_success(True)
+            context.set_exception(fse)
+            self._failover_writer_success_counter.inc()
+            raise fse
+        except Exception as ex:
+            context.set_success(False)
+            context.set_exception(ex)
+            self._failover_writer_failed_counter.inc()
+            raise ex
+        finally:
+            context.close_context()
+            if self._telemetry_failover_additional_top_trace_setting:
+                telemetry_factory.post_copy(context, TelemetryTraceLevel.FORCE_TOP_LEVEL)
 
     def restore_session_state(self, conn: Optional[Connection]):
         """
@@ -401,7 +464,8 @@ def _connect_to(self, host: HostInfo):
             logger.debug("FailoverPlugin.EstablishedConnection", host)
         except Exception as ex:
             if self._plugin_service is not None:
-                logger.debug("FailoverPlugin.ConnectionToHostFailed", 'writer' if host.role == HostRole.WRITER else 'reader', host.url)
+                logger.debug("FailoverPlugin.ConnectionToHostFailed",
+                             'writer' if host.role == HostRole.WRITER else 'reader', host.url)
             raise ex
 
     def _should_attempt_reader_connection(self) -> bool: