[GitHub] [airflow] sekikn commented on issue #6399: [AIRFLOW-5730] Enable get_pandas_df on PinotDbApiHook

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

sekikn commented on issue #6399: [AIRFLOW-5730] Enable get_pandas_df on PinotDbApiHook
URL: https://github.com/apache/airflow/pull/6399#issuecomment-546126938
 
 
   Updated the PR. The CI failure seems to be a transient one. The newly added test works on my local environment, as follows:
   
   ```
   $ git diff upstream/master 
   diff --git a/airflow/contrib/hooks/pinot_hook.py b/airflow/contrib/hooks/pinot_hook.py
   index e617f8e9b..0864b3584 100644
   --- a/airflow/contrib/hooks/pinot_hook.py
   +++ b/airflow/contrib/hooks/pinot_hook.py
   @@ -90,8 +90,5 @@ class PinotDbApiHook(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/contrib/hooks/test_pinot_hook.py b/tests/contrib/hooks/test_pinot_hook.py
   index 72cee01ea..489af15aa 100644
   --- a/tests/contrib/hooks/test_pinot_hook.py
   +++ b/tests/contrib/hooks/test_pinot_hook.py
   @@ -34,6 +34,7 @@ class TestPinotDbApiHook(unittest.TestCase):
            self.conn.conn_type = 'http'
            self.conn.extra_dejson = {'endpoint': 'pql'}
            self.cur = mock.MagicMock()
   +        self.conn.cursor.return_value = self.cur
            self.conn.__enter__.return_value = self.cur
            self.conn.__exit__.return_value = None
    
   @@ -74,3 +75,14 @@ class TestPinotDbApiHook(unittest.TestCase):
            result_sets = [('row1',), ('row2',)]
            self.cur.fetchone.return_value = result_sets[0]
            self.assertEqual(result_sets[0], self.db_hook().get_first(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])
   $ ./run-tests tests.contrib.hooks.test_pinot_hook:TestPinotDbApiHook.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.contrib.hooks.test_pinot_hook:TestPinotDbApiHook.test_get_pandas_df
   
   .
   ----------------------------------------------------------------------
   Ran 1 test in 0.289s
   
   OK
   [2019-10-25 07:20:36,090] {settings.py:200} DEBUG - Disposing DB connection pool (PID 3587)
   ```

----------------------------------------------------------------
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