commit 8f1a585b58e6d8091f4524e6cfb09c606e828825
author Kengo Seki <sekikn@apache.org> Sat Nov 16 08:53:06 2019 +0900
committer Kaxil Naik <kaxilnaik@gmail.com> Fri Nov 15 23:53:06 2019 +0000
tree c487f9f683e800dda76ddab3f28dc68d3528fecb
parent cab6f63015da206128675d3d73668e0d6f093491

[AIRFLOW-5730] Enable get_pandas_df on PinotDbApiHook (#6399)

airflow/contrib/hooks/pinot_hook.py

diff --git a/airflow/contrib/hooks/pinot_hook.py b/airflow/contrib/hooks/pinot_hook.py
index 084722e..c4edc24 100644
--- a/airflow/contrib/hooks/pinot_hook.py
+++ b/airflow/contrib/hooks/pinot_hook.py
@@ -278,8 +278,5 @@
     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()

tests/contrib/hooks/test_pinot_hook.py

diff --git a/tests/contrib/hooks/test_pinot_hook.py b/tests/contrib/hooks/test_pinot_hook.py
index 6229894..3765f7d 100644
--- a/tests/contrib/hooks/test_pinot_hook.py
+++ b/tests/contrib/hooks/test_pinot_hook.py
@@ -180,6 +180,7 @@
         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
 
@@ -220,3 +221,14 @@
         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])