Aws connectors

numalogic/connectors/_config.py

@@ -7,6 +7,7 @@ class ConnectorType(IntEnum):
     redis = 0
     prometheus = 1
     druid = 2
+    rds = 3
 
 
 @dataclass

numalogic/connectors/rds/__init__.py

@@ -0,0 +1,3 @@
+from numalogic.connectors.rds._rds import RDSFetcher
+
+__all__ = ["RDSFetcher"]

numalogic/connectors/rds/_base.py

@@ -0,0 +1,159 @@
+from abc import ABCMeta, abstractmethod
+from typing import Optional
+import pandas as pd
+from numalogic.connectors.utils.aws.config import DatabaseServiceProvider, RDSConfig
+from numalogic.connectors.utils.aws.boto3_client_manager import Boto3ClientManager
+import logging
+from numalogic.connectors._config import Pivot
+import time
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def format_dataframe(
+    df: pd.DataFrame,
+    query: str,
+    datetime_field_name: str,
+    group_by: Optional[list[str]] = None,
+    pivot: Optional[Pivot] = None,
+) -> pd.DataFrame:
+    """
+    Executes formatting operations on a pandas DataFrame.
+
+    Arguments
+    ----------
+    df : pd.DataFrame
+        The input DataFrame to be formatted.
+    query : str
+        The SQL query used to retrieve the data.
+    datetime_field_name : str
+        The name of the datetime field in the DataFrame.
+    group_by : Optional[list[str]], optional
+        A list of column names to group the DataFrame by, by default None.
+    pivot : Optional[Pivot], optional
+        An optional Pivot object specifying the index, columns,
+        and values for pivoting the DataFrame, by default None.
+
+    Returns
+    -------
+    pd.DataFrame : The formatted DataFrame.
+
+    """
+    _start_time = time.perf_counter()
+    df["timestamp"] = pd.to_datetime(df[datetime_field_name]).astype("int64") // 10**6
+    df.drop(columns=datetime_field_name, inplace=True)
+    if group_by:
+        df = df.groupby(by=group_by).sum().reset_index()
+
+    if pivot and pivot.columns:
+        df = df.pivot(
+            index=pivot.index,
+            columns=pivot.columns,
+            values=pivot.value,
+        )
+        df.columns = df.columns.map("{0[1]}".format)
+        df.reset_index(inplace=True)
+    _end_time = time.perf_counter() - _start_time
+    _LOGGER.info("RDS MYSQL Query: %s, Format time:  %.4fs", query, _end_time)
+    return df
+
+
+class RDSBase(metaclass=ABCMeta):
+    """
+    class represents a data fetcher for RDS (Relational Database Service) connections. It
+    provides methods for retrieving the RDS token, getting the password, establishing a
+    connection, and executing queries.
+
+    Args:
+        - db_config (RDSConfig): The configuration object for the RDS connection.
+        - kwargs (dict): Additional keyword arguments.
+
+    """
+
+    def __init__(self, db_config: RDSConfig, **kwargs):
+        self.kwargs = kwargs
+        self.db_config = db_config
+        self.connection = None
+        self.database_type = db_config.database_type
+        self.boto3_client_manager = Boto3ClientManager(self.db_config)
+
+    def get_rds_token(self) -> str:
+        """
+        Generates an RDS authentication token using the provided RDS boto3 client.
+
+        Args:
+            - rds_boto3_client (boto3.client): The RDS boto3 client used to generate the
+            authentication token.
+
+        Returns
+        -------
+            str: The generated RDS authentication token.
+
+        """
+        rds_client = self.boto3_client_manager.get_client(DatabaseServiceProvider.RDS)
+        return self.boto3_client_manager.get_rds_token(rds_client)
+
+    def get_password(self) -> str:
+        """
+        Retrieves the password for the RDS connection.
+
+        If 'aws_rds_use_iam' is True, it calls the get_rds_token() method to generate the RDS
+        token. Otherwise, it returns the database password from the configuration.
+
+        Returns
+        -------
+            str: The password for the RDS connection.
+
+        """
+        if self.db_config.aws_rds_use_iam:
+            _LOGGER.info("using aws_rds_use_iam to generate RDS Token")
+            return self.get_rds_token()
+
+        _LOGGER.info("using password from config to connect RDS Database")
+        return self.db_config.database_password
+
+    @abstractmethod
+    def get_connection(self):
+        """
+        Establishes a connection to the RDS database.
+
+        This method is a placeholder and needs to be implemented in a subclass. It should handle
+        the logic for establishing a connection to the RDS database based on the provided
+        configuration.
+
+        Returns
+        -------
+            None
+
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_db_cursor(self, *args, **kwargs):
+        """
+        Retrieves a database cursor for executing queries.
+
+        This method is a placeholder and needs to be implemented in a subclass. It should handle
+        the logic for retrieving a database cursor based on the established connection.
+
+        Returns
+        -------
+            None
+
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def execute_query(self, query: str) -> pd.DataFrame:
+        """
+        Executes a query on the RDS database and returns the result as a pandas DataFrame.
+
+        Args:
+                query (str): The SQL query to be executed.
+
+        Returns
+        -------
+            pd.DataFrame: The result of the query as a pandas DataFrame.
+
+        """
+        raise NotImplementedError

numalogic/connectors/rds/_rds.py

@@ -0,0 +1,70 @@
+from typing import Optional
+from numalogic.connectors._base import DataFetcher
+from numalogic.connectors._config import Pivot
+from numalogic.connectors.rds._base import format_dataframe
+from numalogic.connectors.utils.aws.config import RDSConfig
+import logging
+import pandas as pd
+from numalogic.connectors.rds.db.factory import RdsFactory
+import time
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class RDSFetcher(DataFetcher):
+    """
+    class is a subclass of DataFetcher and ABC (Abstract Base Class).
+    It is used to fetch data from an RDS (Relational Database Service) instance by executing
+    a given SQL query.
+
+    Attributes
+    ----------
+        db_config (RDSConfig): The configuration object for the RDS instance.
+        fetcher (db.CLASS_TYPE): The fetcher object for the specific database type.
+
+    """
+
+    def __init__(self, db_config: RDSConfig):
+        super().__init__(db_config.endpoint)
+        self.db_config = db_config
+        factory_object = RdsFactory()
+        self.fetcher = factory_object.get_db_handler(db_config.database_type.lower())(db_config)
+        _LOGGER.info("Executing for database type: %s", self.fetcher.database_type)
+
+    def fetch(
+        self,
+        query,
+        datetime_field_name: str,
+        pivot: Optional[Pivot] = None,
+        group_by: Optional[list[str]] = None,
+    ) -> pd.DataFrame:
+        """
+        Fetches data from the RDS instance by executing the given query.
+
+        Args:
+            query (str): The SQL query to be executed.
+            datetime_field_name (str): The name of the datetime field in the fetched data.
+            pivot (Optional[Pivot], optional): The pivot configuration for the fetched data.
+            Defaults to None.
+            group_by (Optional[list[str]], optional): The list of fields to group the
+            fetched data by. Defaults to None.
+
+        Returns
+        -------
+            pd.DataFrame: A pandas DataFrame containing the fetched data.
+        """
+        _start_time = time.perf_counter()
+        df = self.fetcher.execute_query(query)
+        if df.empty or df.shape[0] == 0:
+            _LOGGER.warning("No data found for query : %s ", query)
+            return pd.DataFrame()
+
+        formatted_df = format_dataframe(
+            df, query=query, datetime_field_name=datetime_field_name, pivot=pivot, group_by=group_by
+        )
+        _end_time = time.perf_counter() - _start_time
+        _LOGGER.info("RDS Query: %s Fetch Time: %.4fs", query, _end_time)
+        return formatted_df
+
+    def raw_fetch(self, *args, **kwargs) -> pd.DataFrame:
+        raise NotImplementedError

numalogic/connectors/rds/db/factory.py

@@ -0,0 +1,34 @@
+import logging
+
+from numalogic.connectors.utils.aws.config import DatabaseTypes
+from numalogic.connectors.utils.aws.exceptions import UnRecognizedDatabaseTypeException
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class RdsFactory:
+    """class represents a factory for creating database handlers for different database types."""
+
+    @classmethod
+    def get_db_handler(cls, database_type: DatabaseTypes):
+        """
+        Get the database handler for the specified database type.
+
+        Args:
+            - database_type (str): The type of the database.
+
+        Returns
+        -------
+            - The database handler for the specified database type.
+
+        Raises
+        ------
+            - UnRecognizedDatabaseTypeException: If the specified database type is not supported.
+
+        """
+        if database_type == DatabaseTypes.MYSQL:
+            from numalogic.connectors.rds.db.mysql_fetcher import MysqlFetcher
+
+            return MysqlFetcher
+
+        raise UnRecognizedDatabaseTypeException(f"database_type: {database_type} is not supported")