Commit eb60edcb authored by Huihui, Jonathan's avatar Huihui, Jonathan
Browse files

Merge branch 'feature/fix_query_pandas' into 'develop' 

Feature/fix query pandas

See merge request nset-utilities/common-package!20
parents acaab61a 59b5f728

src/common/mixins/pandas.py

+19 −6
Original line number Diff line number Diff line
@@ -26,12 +26,25 @@ class PandasMixin(): 
                self.logger.error('Cursor is not open, please create one.')

            # TODO wrap in a try except block

            self.logger.debug(f'Returning {query} as pandas dataframe.')

            data_fetch_query = self.cursor.execute(query)

            if data_fetch_query is None:  # FIXME better way to handle this?

                self.logger.error(f"Unable to fetch {query}.")

                # FIXME raise an Exception here?

            dataframe = pd.DataFrame(data_fetch_query.fetchall())

            dataframe.columns = data_fetch_query.keys()



            self.cursor.execute(query)

            # https://www.psycopg.org/docs/cursor.html

            # Comments on the execute method:

            #   The method returns None. If a query was executed,

            #   the returned values can be retrieved using fetch*() methods.

            result = self.cursor.fetchall()

            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}'")



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



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

            # dataframe.columns = result.keys()

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

            return dataframe



        if query is not None and not isinstance(self.connection,

src/requirements.txt

+1 −0
Original line number Diff line number Diff line
@@ -2,3 +2,4 @@ rich==10.11.0 
pandas==1.3.3

psycopg2-binary==2.9.1

pymssql==2.2.7

SQLAlchemy~=1.4.44

src/setup.py

+5 −3
Original line number Diff line number Diff line
@@ -15,20 +15,22 @@ _here = os.path.abspath(os.path.dirname(__file__)) 
_requirements_file = _here + '/requirements.txt'

_requirements = []

if os.path.isfile(_requirements_file):

    with open(_requirements_file, 'r') as f:

    with open(_requirements_file, 'r', encoding='utf-8') as f:

        _requirements = f.read().splitlines()

_license = 'MIT'

_license_file = _here + 'LICENSE'

if os.path.isfile(_license_file):

    with open(_license_file, 'r') as f:

    with open(_license_file, 'r', encoding='utf-8') as f:

        _license = f.read()

_description_file = _here + 'README.md'

_long_description = ''

if os.path.isfile(_description_file):

    with open(_description_file, 'r') as f:

    with open(_description_file, 'r', encoding='utf-8') as f:

        _long_description = f.read()

_scripts_dir = 'scripts/'

_scripts = os.listdir(_here + '/scripts/')

if '.connections' in _scripts:

    _scripts.remove('.connections')

_scripts = ['scripts/' + script for script in _scripts]



# setup

test/test.py

+15 −0
Original line number Diff line number Diff line
@@ -10,6 +10,7 @@ from datetime import datetime 
from webbrowser import get

from common.database import Database

from common.mixins.postgres import PostgresMixin

from common.mixins.pandas import PandasMixin

from common.mixins.mssql import MSSQLMixin

from common.mixins.postgres_logger import PostgresLoggerMixin

from common.logz import create_logger
@@ -298,5 +299,19 @@ class MailTestCase(unittest.TestCase): 
        #            recipients='huihuijk@ornl.gov, jonathanhuihui@yahoo.com, jonathankhuihui@gmail.com')





class PandasTestCase(unittest.TestCase):

    class PandasDB(Database, PostgresMixin, PandasMixin):

        """Class to query postgresdb with pandas"""



    def test_query_pandas(self):

        with self.PandasDB() as db:

            # Query the database to get a dataframe

            df = db.query_pandas(query="SELECT * FROM public.test_table",

                                 table="test_table", schema="public")

            # Ensure the dataframe as expected values

            self.assertTrue(1 in df['id'].values)

            self.assertTrue("hello_world" in df['test_string'].values)





if __name__ == "__main__":

    unittest.main(verbosity=2)