Commit 20fd1076 authored by Grant's avatar Grant
Browse files

enforce optional logging in PandasMixin

parent 513f045f

COMMON_VERSION

+1 −1
Original line number Diff line number Diff line 
0.3.3
 
0.3.4

src/common/mixins/pandas_mixin.py

+7 −3
Original line number Diff line number Diff line
@@ -8,7 +8,7 @@ import sqlalchemy 
class PandasMixin:

    """Group together all pandas logic."""



    def query_pandas(self, query=None, table=None, schema=None):

    def query_pandas(self, query=None, table=None, schema=None, log_df=False, log_query=False):

        """Return a query as a pandas table.



        Provided a query, uses the default cursor assumming a
@@ -17,6 +17,8 @@ class PandasMixin: 
        :param query: a query string to return as a table

        :param table: the table name to retrieve data from

        :param schema: the schema to select the table from

        :param log_df: should the dataframe be logged? defaults to False

        :param log_query: should the query be logged? defaults to False

        :returns: a pandas.DataFframe object

        """
@@ -37,12 +39,14 @@ class PandasMixin: 
            if result is None:

                self.logger.error("Query returned: 'None' for query " f"'{query}'")

            else:

                self.logger.info(f"Query returned: '{result}' for query " f"'{query}'")

                if log_query:

                    self.logger.debug(f"Query returned: '{result}' for query " f"'{query}'")



            columns = [i[0] for i in self.cursor.description]



            dataframe = pd.DataFrame.from_records(result, columns=columns)

            # dataframe.columns = result.keys()

            if log_df:

                self.logger.info(f"Fetched Dataframe: {dataframe}")

            return dataframe