boto3-assist 0.32.0__py3-none-any.whl

boto3_assist/connection.py

@@ -0,0 +1,146 @@
+"""
+Geek Cafe, LLC
+Maintainers: Eric Wilson
+MIT License.  See Project Root for the license information.
+"""
+
+from typing import Optional, List
+
+from aws_lambda_powertools import Logger
+from boto3_assist.boto3session import Boto3SessionManager
+from boto3_assist.environment_services.environment_variables import (
+    EnvironmentVariables,
+)
+from boto3_assist.connection_tracker import ConnectionTracker
+
+
+logger = Logger()
+tracker: ConnectionTracker = ConnectionTracker()
+
+
+class Connection:
+    """Base Boto 3 Connection"""
+
+    def __init__(
+        self,
+        *,
+        service_name: Optional[str] = None,
+        aws_profile: Optional[str] = None,
+        aws_region: Optional[str] = None,
+        aws_access_key_id: Optional[str] = None,
+        aws_secret_access_key: Optional[str] = None,
+        aws_end_point_url: Optional[str] = None,
+        assume_role_arn: Optional[str] = None,
+        assume_role_chain: Optional[List[str]] = None,
+        assume_role_duration_seconds: Optional[int] = 3600,
+    ) -> None:
+        self.__aws_profile = aws_profile
+        self.__aws_region = aws_region
+        self.__aws_access_key_id = aws_access_key_id
+        self.__aws_secret_access_key = aws_secret_access_key
+        self.end_point_url = aws_end_point_url
+        self.__session: Boto3SessionManager | None = None
+        self.__assume_role_arn: Optional[str] = assume_role_arn
+        self.__service_name: str | None = service_name
+        self.__assume_role_chain = assume_role_chain
+        self.__assume_role_duration_seconds = assume_role_duration_seconds
+        if self.__service_name is None:
+            raise RuntimeError(
+                "Service Name is not available. The service name is required."
+            )
+
+        self.raise_on_error: bool = True
+
+    def setup(self, setup_source: Optional[str] = None) -> None:
+        """
+        Setup the environment.  Automatically called via init.
+        You can run setup at anytime with new parameters.
+        Args: setup_source: Optional[str] = None
+            Defines the source of the setup.  Useful for logging.
+        Returns: None
+        """
+
+        logger.debug(
+            {
+                "metric_filter": f"{self.service_name}_connection_setup",
+                "source": f"{self.service_name} Connection",
+                "aws_profile": self.aws_profile,
+                "aws_region": self.aws_region,
+                "setup_source": setup_source,
+            }
+        )
+
+        self.__session = Boto3SessionManager(
+            service_name=self.service_name,
+            aws_profile=self.aws_profile,
+            aws_region=self.aws_region,
+            aws_access_key_id=self.aws_access_key_id,
+            aws_secret_access_key=self.aws_secret_access_key,
+            aws_endpoint_url=self.end_point_url,
+            assume_role_arn=self.__assume_role_arn,
+            assume_role_chain=self.__assume_role_chain,
+            assume_role_duration_seconds=self.__assume_role_duration_seconds,
+        )
+
+        tracker.add(service_name=self.service_name)
+
+    @property
+    def asw_profile(self) -> str | None:
+        """The AWS Profile"""
+        return self.__aws_profile or EnvironmentVariables.AWS.profile()
+
+    @asw_profile.setter
+    def aws_profile(self, value: str | None):
+        self.__aws_profile = value
+
+    @property
+    def aws_region(self) -> str | None:
+        """The AWS Region"""
+        return self.__aws_region or EnvironmentVariables.AWS.region()
+
+    @aws_region.setter
+    def aws_region(self, value: str | None):
+        self.__aws_region = value
+
+    @property
+    def aws_access_key_id(self) -> str | None:
+        """The AWS Access Key"""
+        return self.__aws_access_key_id or EnvironmentVariables.AWS.aws_access_key_id()
+
+    @aws_access_key_id.setter
+    def aws_access_key_id(self, value: str | None):
+        self.__aws_access_key_id = value
+
+    @property
+    def aws_secret_access_key(self) -> str | None:
+        """The AWS Access Key"""
+        return (
+            self.__aws_secret_access_key
+            or EnvironmentVariables.AWS.aws_secret_access_key()
+        )
+
+    @aws_secret_access_key.setter
+    def aws_secret_access_key(self, value: str | None):
+        self.__aws_secret_access_key = value
+
+    @property
+    def service_name(self) -> str:
+        """Service Name"""
+        if self.__service_name is None:
+            raise RuntimeError("Service Name is not available")
+        return self.__service_name
+
+    @service_name.setter
+    def service_name(self, value: str):
+        logger.debug("Setting Service Name")
+        self.__service_name = value
+
+    @property
+    def session(self) -> Boto3SessionManager:
+        """Session"""
+        if self.__session is None:
+            self.setup(setup_source="session init")
+
+        if self.__session is None:
+            raise RuntimeError("Session is not available")
+        return self.__session

boto3_assist/connection_tracker.py

@@ -0,0 +1,120 @@
+"""
+Geek Cafe, LLC
+Maintainers: Eric Wilson
+MIT License. See Project Root for the license information.
+"""
+
+import traceback
+import os
+from typing import Dict
+
+
+class ConnectionTracker:
+    """
+    Tracks Boto3 connection requests for performance tuning and debugging.
+
+    Attributes:
+        __stack_trace_env_var (str): Environment variable name to enable stack trace logging.
+        __issue_stack_trace (bool | None): Caches the result of whether stack trace logging is enabled.
+        __connection_counter (Dict[str, int]): Tracks the number of connections per service.
+    """
+
+    def __init__(self) -> None:
+        self.__stack_trace_env_var: str = "BOTO3_ASSIST_CONNECTION_STACK_TRACE"
+        self.__issue_stack_trace: bool | None = None
+        self.__connection_counter: Dict[str, int] = {}
+
+    def add(self, service_name: str) -> None:
+        """
+        Increments the connection count for a given service and
+        performs a check on the number of connections.
+
+        Args:
+            service_name (str): Name of the AWS service.
+        """
+        self.__connection_counter[service_name] = (
+            self.__connection_counter.get(service_name, 0) + 1
+        )
+
+        self.check(service_name=service_name)
+
+    @property
+    def issue_stack_trace(self) -> bool:
+        """
+        Checks if stack trace logging is enabled by the environment variable.
+
+        Returns:
+            bool: True if stack trace logging is enabled, False otherwise.
+        """
+        if self.__issue_stack_trace is None:
+            self.__issue_stack_trace = (
+                os.getenv(self.__stack_trace_env_var, "").lower() == "true"
+            )
+        return self.__issue_stack_trace
+
+    def check(self, service_name: str) -> None:
+        """
+        Checks the connection count for a service and logs warnings if needed.
+
+        Args:
+            service_name (str): Name of the AWS service.
+        """
+        connection_count = self.__connection_counter.get(service_name, 0)
+        if connection_count > 1:
+            service_message = (
+                f"Your {service_name} service has more than one connection.\n"
+            )
+
+            if not self.issue_stack_trace:
+                stack_trace_message = (
+                    f"📄 NOTE: To add additional information 👀 to the log and determine where additional connections are being created: "
+                    f"set the environment variable 👉{self.__stack_trace_env_var}👈 to true ✅. \n"
+                )
+            else:
+                stack = "\n".join(traceback.format_stack())
+                stack_trace_message = (
+                    f"\nStack Trace Enabled with {self.__stack_trace_env_var}\n{stack}"
+                )
+
+            self.__log_warning(
+                f"{service_message}"
+                f"Your boto3 connection count is {connection_count}. "
+                "Under most circumstances, you should be able to use the same connection "
+                "instead of creating a new one. Connections are expensive in terms of time and latency. "
+                "If you are seeing performance issues, check how and where you are creating your "
+                "connections. You should be able to pass the connection to your other objects "
+                "and reuse your boto3 connections. "
+                "\n🧪 MOCK Testing may show this message as well, in which case you can dismiss this warning.🧪\n"
+                f"{stack_trace_message}"
+            )
+
+    def decrement_connection(self, service_name: str) -> None:
+        """
+        Decrements the connection count for a service.
+
+        Args:
+            service_name (str): Name of the AWS service.
+        """
+        if (
+            service_name in self.__connection_counter
+            and self.__connection_counter[service_name] > 0
+        ):
+            self.__connection_counter[service_name] -= 1
+
+    def reset(self, service_name: str) -> None:
+        """
+        Resets the connection count for a service to zero.
+
+        Args:
+            service_name (str): Name of the AWS service.
+        """
+        self.__connection_counter[service_name] = 0
+
+    def __log_warning(self, message: str) -> None:
+        """
+        Logs a warning message.
+
+        Args:
+            message (str): The warning message to log.
+        """
+        print(f"Warning: {message}")