[GitHub] [airflow] sekikn commented on issue #6057: [AIRFLOW-5442] implementing get_pandas_df method for druid broker hook

[GitBox <gi...@apache.org>]

sekikn commented on issue #6057: [AIRFLOW-5442] implementing get_pandas_df method for druid broker hook
URL: https://github.com/apache/airflow/pull/6057#issuecomment-545940951
 
 
   @blcksrx How about adding the following test to tests/hooks/test_druid_hook.py, just as for `get_records` and `get_first`?
   
   ```
       def test_get_pandas_df(self):
           statement = 'SQL'
           column = 'col'
           result_sets = [('row1',), ('row2',)]
           self.cur.description = [(column,)]
           self.cur.fetchall.return_value = result_sets
           df = self.db_hook().get_pandas_df(statement)
   
           self.assertEqual(column, df.columns[0])
           for i in range(len(result_sets)):  # pylint: disable=consider-using-enumerate
               self.assertEqual(result_sets[i][0], df.values.tolist()[i][0])
           assert self.conn.close.call_count == 1
           assert self.cur.close.call_count == 1
           self.cur.execute.assert_called_once_with(statement)
   ```
   
   This should work, as follows:
   
   ```
   $ git diff
   diff --git a/airflow/hooks/druid_hook.py b/airflow/hooks/druid_hook.py
   index c3cd3cd71..e2e20f1ec 100644
   --- a/airflow/hooks/druid_hook.py
   +++ b/airflow/hooks/druid_hook.py
   @@ -158,8 +158,5 @@ class DruidDbApiHook(DbApiHook):
        def set_autocommit(self, conn, autocommit):
            raise NotImplementedError()
    
   -    def get_pandas_df(self, sql, parameters=None):
   -        raise NotImplementedError()
   -
        def insert_rows(self, table, rows, target_fields=None, commit_every=1000):
            raise NotImplementedError()
   diff --git a/tests/hooks/test_druid_hook.py b/tests/hooks/test_druid_hook.py
   index 867757a10..a20e8e42f 100644
   --- a/tests/hooks/test_druid_hook.py
   +++ b/tests/hooks/test_druid_hook.py
   @@ -181,6 +181,21 @@ class TestDruidDbApiHook(unittest.TestCase):
            assert self.cur.close.call_count == 1
            self.cur.execute.assert_called_once_with(statement)
    
   +    def test_get_pandas_df(self):
   +        statement = 'SQL'
   +        column = 'col'
   +        result_sets = [('row1',), ('row2',)]
   +        self.cur.description = [(column,)]
   +        self.cur.fetchall.return_value = result_sets
   +        df = self.db_hook().get_pandas_df(statement)
   +
   +        self.assertEqual(column, df.columns[0])
   +        for i in range(len(result_sets)):  # pylint: disable=consider-using-enumerate
   +            self.assertEqual(result_sets[i][0], df.values.tolist()[i][0])
   +        assert self.conn.close.call_count == 1
   +        assert self.cur.close.call_count == 1
   +        self.cur.execute.assert_called_once_with(statement)
   +
    
    if __name__ == '__main__':
        unittest.main()
   $ ./run-tests tests.hooks.test_druid_hook:TestDruidDbApiHook.test_get_pandas_df
   AIRFLOW__CORE__SQL_ALCHEMY_CONN not set - using default
   Airflow home: /home/sekikn
   Airflow root: /home/sekikn/repos/airflow
   Home of the user: /home/sekikn
   
   
   Skipping initializing of the DB as it was initialized already
   
   You can re-initialize the database by adding --with-db-init flag when running tests
   
   KRB5_KTNAME variable is empty - no kerberos intialisation
   
   Starting the tests with arguments: tests.hooks.test_druid_hook:TestDruidDbApiHook.test_get_pandas_df
   
   .
   ----------------------------------------------------------------------
   Ran 1 test in 0.177s
   
   OK
   ```

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services