[v3-0-test] test: migrate get_pandas_df to get_df in provider t…

airflow-core/docs/extra-packages-ref.rst

@@ -106,6 +106,8 @@ python dependencies for the provided package.
 +---------------------+-----------------------------------------------------+----------------------------------------------------------------------------+
 | pandas              | ``pip install 'apache-airflow[pandas]'``            | Install Pandas library compatible with Airflow                             |
 +---------------------+-----------------------------------------------------+----------------------------------------------------------------------------+
+| polars              | ``pip install 'apache-airflow[polars]'``            | Polars hooks and operators                                                 |
++---------------------+-----------------------------------------------------+----------------------------------------------------------------------------+
 | rabbitmq            | ``pip install 'apache-airflow[rabbitmq]'``          | RabbitMQ support as a Celery backend                                       |
 +---------------------+-----------------------------------------------------+----------------------------------------------------------------------------+
 | sentry              | ``pip install 'apache-airflow[sentry]'``            | Sentry service for application logging and monitoring                      |

airflow-core/pyproject.toml

@@ -138,7 +138,7 @@ dependencies = [
     # pre-installed providers
     "apache-airflow-providers-common-compat>=1.6.0",
     "apache-airflow-providers-common-io>=1.5.3",
-    "apache-airflow-providers-common-sql>=1.25.0",
+    "apache-airflow-providers-common-sql>=1.26.0",
     "apache-airflow-providers-smtp>=2.0.2",
     "apache-airflow-providers-standard>=0.4.0",
 ]
@@ -237,7 +237,7 @@ dev = [
     "apache-airflow-task-sdk",
     # TODO(potiuk): eventually we do not want any providers nor apache-airflow extras to be needed for
     # airflow-core tests
-    "apache-airflow[pandas]",
+    "apache-airflow[pandas,polars]",
     "apache-airflow-providers-amazon",
     "apache-airflow-providers-celery",
     "apache-airflow-providers-cncf-kubernetes",

providers/apache/drill/README.rst

@@ -54,7 +54,7 @@ Requirements
 PIP package                              Version required
 =======================================  ===========================
 ``apache-airflow``                       ``>=2.9.0``
-``apache-airflow-providers-common-sql``  ``>=1.20.0``
+``apache-airflow-providers-common-sql``  ``>=1.26.0``
 ``sqlalchemy-drill``                     ``>=1.1.0,!=1.1.6,!=1.1.7``
 =======================================  ===========================
 

providers/apache/drill/pyproject.toml

@@ -58,7 +58,7 @@ requires-python = "~=3.9"
 # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
 dependencies = [
     "apache-airflow>=2.9.0",
-    "apache-airflow-providers-common-sql>=1.20.0",
+    "apache-airflow-providers-common-sql>=1.26.0",
     # Workaround until we get https://github.com/JohnOmernik/sqlalchemy-drill/issues/94 fixed.
     "sqlalchemy-drill>=1.1.0,!=1.1.6,!=1.1.7",
 ]
@@ -70,7 +70,7 @@ dev = [
     "apache-airflow-devel-common",
     "apache-airflow-providers-common-sql",
     # Additional devel dependencies (do not remove this line and add extra development dependencies)
-    "apache-airflow-providers-common-sql[pandas]",
+    "apache-airflow-providers-common-sql[pandas,polars]",
 ]
 
 # To build docs:

providers/apache/drill/tests/unit/apache/drill/hooks/test_drill.py

@@ -90,17 +90,33 @@ def test_get_records(self):
         assert self.cur.close.call_count == 1
         self.cur.execute.assert_called_once_with(statement)
 
-    def test_get_pandas_df(self):
+    def test_get_df_pandas(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)
+        df = self.db_hook().get_df(statement, df_type="pandas")
 
         assert column == df.columns[0]
         for i, item in enumerate(result_sets):
             assert item[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)
+
+    def test_get_df_polars(self):
+        statement = "SQL"
+        column = "col"
+        result_sets = [("row1",), ("row2",)]
+        mock_execute = MagicMock()
+        mock_execute.description = [(column, None, None, None, None, None, None)]
+        mock_execute.fetchall.return_value = result_sets
+        self.cur.execute.return_value = mock_execute
+        df = self.db_hook().get_df(statement, df_type="polars")
+
+        self.cur.execute.assert_called_once_with(statement)
+        mock_execute.fetchall.assert_called_once_with()
+        assert column == df.columns[0]
+        assert result_sets[0][0] == df.row(0)[0]
+        assert result_sets[1][0] == df.row(1)[0]

providers/apache/druid/README.rst

@@ -54,7 +54,7 @@ Requirements
 PIP package                              Version required
 =======================================  ==================
 ``apache-airflow``                       ``>=2.9.0``
-``apache-airflow-providers-common-sql``  ``>=1.20.0``
+``apache-airflow-providers-common-sql``  ``>=1.26.0``
 ``pydruid``                              ``>=0.4.1``
 =======================================  ==================
 

providers/apache/druid/pyproject.toml

@@ -58,7 +58,7 @@ requires-python = "~=3.9"
 # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
 dependencies = [
     "apache-airflow>=2.9.0",
-    "apache-airflow-providers-common-sql>=1.20.0",
+    "apache-airflow-providers-common-sql>=1.26.0",
     "pydruid>=0.4.1",
 ]
 
@@ -77,6 +77,7 @@ dev = [
     "apache-airflow-providers-apache-hive",
     "apache-airflow-providers-common-sql",
     # Additional devel dependencies (do not remove this line and add extra development dependencies)
+    "apache-airflow-providers-common-sql[polars]",
 ]
 
 # To build docs:

providers/apache/druid/tests/unit/apache/druid/hooks/test_druid.py

@@ -452,17 +452,31 @@ def test_get_records(self):
         assert self.cur.close.call_count == 1
         self.cur.execute.assert_called_once_with(statement)
 
-    def test_get_pandas_df(self):
+    def test_get_df_pandas(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)
+        df = self.db_hook().get_df(statement, df_type="pandas")
 
         assert column == df.columns[0]
         for i, item in enumerate(result_sets):
             assert item[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)
+
+    def test_get_df_polars(self):
+        statement = "SQL"
+        column = "col"
+        result_sets = [("row1",), ("row2",)]
+        mock_execute = MagicMock()
+        mock_execute.description = [(column, None, None, None, None, None, None)]
+        mock_execute.fetchall.return_value = result_sets
+        self.cur.execute.return_value = mock_execute
+
+        df = self.db_hook().get_df(statement, df_type="polars")
+        assert column == df.columns[0]
+        assert result_sets[0][0] == df.row(0)[0]
+        assert result_sets[1][0] == df.row(1)[0]

providers/apache/impala/README.rst

@@ -54,7 +54,7 @@ Requirements
 PIP package                              Version required
 =======================================  ==================
 ``impyla``                               ``>=0.18.0,<1.0``
-``apache-airflow-providers-common-sql``  ``>=1.20.0``
+``apache-airflow-providers-common-sql``  ``>=1.26.0``
 ``apache-airflow``                       ``>=2.9.0``
 =======================================  ==================
 

providers/apache/impala/pyproject.toml

@@ -58,7 +58,7 @@ requires-python = "~=3.9"
 # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
 dependencies = [
     "impyla>=0.18.0,<1.0",
-    "apache-airflow-providers-common-sql>=1.20.0",
+    "apache-airflow-providers-common-sql>=1.26.0",
     "apache-airflow>=2.9.0",
 ]
 
@@ -77,7 +77,7 @@ dev = [
     "apache-airflow-providers-common-sql",
     # Additional devel dependencies (do not remove this line and add extra development dependencies)
     "kerberos>=1.3.0",
-    "apache-airflow-providers-common-sql[pandas]"
+    "apache-airflow-providers-common-sql[pandas,polars]"
 ]
 
 # To build docs:

providers/apache/impala/tests/unit/apache/impala/hooks/test_impala.py

@@ -107,17 +107,32 @@ def test_get_records(impala_hook_fixture):
     impala_hook_fixture.get_conn.return_value.cursor.return_value.execute.assert_called_once_with(statement)
 
 
-def test_get_pandas_df(impala_hook_fixture):
+def test_get_df(impala_hook_fixture):
     statement = "SQL"
     column = "col"
     result_sets = [("row1",), ("row2",)]
     impala_hook_fixture.get_conn.return_value.cursor.return_value.description = [(column,)]
     impala_hook_fixture.get_conn.return_value.cursor.return_value.fetchall.return_value = result_sets
-    df = impala_hook_fixture.get_pandas_df(statement)
+    df = impala_hook_fixture.get_df(statement, df_type="pandas")
 
     assert column == df.columns[0]
 
     assert result_sets[0][0] == df.values.tolist()[0][0]
     assert result_sets[1][0] == df.values.tolist()[1][0]
 
     impala_hook_fixture.get_conn.return_value.cursor.return_value.execute.assert_called_once_with(statement)
+
+
+def test_get_df_polars(impala_hook_fixture):
+    statement = "SQL"
+    column = "col"
+    result_sets = [("row1",), ("row2",)]
+    mock_execute = MagicMock()
+    mock_execute.description = [(column, None, None, None, None, None, None)]
+    mock_execute.fetchall.return_value = result_sets
+    impala_hook_fixture.get_conn.return_value.cursor.return_value.execute.return_value = mock_execute
+
+    df = impala_hook_fixture.get_df(statement, df_type="polars")
+    assert column == df.columns[0]
+    assert result_sets[0][0] == df.row(0)[0]
+    assert result_sets[1][0] == df.row(1)[0]

providers/apache/pinot/README.rst

@@ -54,7 +54,7 @@ Requirements
 PIP package                              Version required
 =======================================  ==================
 ``apache-airflow``                       ``>=2.9.0``
-``apache-airflow-providers-common-sql``  ``>=1.20.0``
+``apache-airflow-providers-common-sql``  ``>=1.26.0``
 ``pinotdb``                              ``>=5.1.0``
 =======================================  ==================
 

providers/apache/pinot/pyproject.toml

@@ -58,7 +58,7 @@ requires-python = "~=3.9"
 # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
 dependencies = [
     "apache-airflow>=2.9.0",
-    "apache-airflow-providers-common-sql>=1.20.0",
+    "apache-airflow-providers-common-sql>=1.26.0",
     "pinotdb>=5.1.0",
 ]
 
@@ -69,7 +69,7 @@ dev = [
     "apache-airflow-devel-common",
     "apache-airflow-providers-common-sql",
     # Additional devel dependencies (do not remove this line and add extra development dependencies)
-    "apache-airflow-providers-common-sql[pandas]"
+    "apache-airflow-providers-common-sql[pandas,polars]"
 ]
 
 # To build docs:

providers/apache/pinot/tests/unit/apache/pinot/hooks/test_pinot.py

@@ -266,17 +266,30 @@ def test_get_first(self):
         self.cur.fetchone.return_value = result_sets[0]
         assert result_sets[0] == self.db_hook().get_first(statement)
 
-    def test_get_pandas_df(self):
+    def test_get_df_pandas(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)
+        df = self.db_hook().get_df(statement, df_type="pandas")
         assert column == df.columns[0]
         for i, item in enumerate(result_sets):
             assert item[0] == df.values.tolist()[i][0]
 
+    def test_get_df_polars(self):
+        statement = "SQL"
+        column = "col"
+        result_sets = [("row1",), ("row2",)]
+        mock_execute = mock.MagicMock()
+        mock_execute.description = [(column, None, None, None, None, None, None)]
+        mock_execute.fetchall.return_value = result_sets
+        self.cur.execute.return_value = mock_execute
+        df = self.db_hook().get_df(statement, df_type="polars")
+        assert column == df.columns[0]
+        assert result_sets[0][0] == df.row(0)[0]
+        assert result_sets[1][0] == df.row(1)[0]
+
 
 class TestPinotAdminHookWithAuth:
     def setup_method(self):

providers/elasticsearch/README.rst

@@ -54,7 +54,7 @@ Requirements
 PIP package                              Version required
 =======================================  ==================
 ``apache-airflow``                       ``>=2.9.0``
-``apache-airflow-providers-common-sql``  ``>=1.20.0``
+``apache-airflow-providers-common-sql``  ``>=1.26.0``
 ``elasticsearch``                        ``>=8.10,<9``
 =======================================  ==================
 

providers/elasticsearch/pyproject.toml

@@ -58,7 +58,7 @@ requires-python = "~=3.9"
 # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
 dependencies = [
     "apache-airflow>=2.9.0",
-    "apache-airflow-providers-common-sql>=1.20.0",
+    "apache-airflow-providers-common-sql>=1.26.0",
     "elasticsearch>=8.10,<9",
 ]
 
@@ -69,7 +69,7 @@ dev = [
     "apache-airflow-devel-common",
     "apache-airflow-providers-common-sql",
     # Additional devel dependencies (do not remove this line and add extra development dependencies)
-    "apache-airflow-providers-common-sql[pandas]",
+    "apache-airflow-providers-common-sql[pandas,polars]",
 ]
 
 # To build docs:

providers/elasticsearch/tests/unit/elasticsearch/hooks/test_elasticsearch.py

@@ -166,9 +166,9 @@ def test_get_records(self):
         self.spy_agency.assert_spy_called(self.cur.close)
         self.spy_agency.assert_spy_called(self.cur.execute)
 
-    def test_get_pandas_df(self):
+    def test_get_df_pandas(self):
         statement = "SELECT * FROM hollywood.actors"
-        df = self.db_hook.get_pandas_df(statement)
+        df = self.db_hook.get_df(statement, df_type="pandas")
 
         assert list(df.columns) == ["index", "name", "firstname", "age"]
         assert df.values.tolist() == ROWS

providers/sqlite/README.rst

@@ -54,7 +54,7 @@ Requirements
 PIP package                              Version required
 =======================================  ==================
 ``apache-airflow``                       ``>=2.9.0``
-``apache-airflow-providers-common-sql``  ``>=1.20.0``
+``apache-airflow-providers-common-sql``  ``>=1.26.0``
 =======================================  ==================
 
 Cross provider package dependencies

providers/sqlite/pyproject.toml

@@ -58,7 +58,7 @@ requires-python = "~=3.9"
 # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
 dependencies = [
     "apache-airflow>=2.9.0",
-    "apache-airflow-providers-common-sql>=1.20.0",
+    "apache-airflow-providers-common-sql>=1.26.0",
 ]
 
 [dependency-groups]
@@ -68,7 +68,7 @@ dev = [
     "apache-airflow-devel-common",
     "apache-airflow-providers-common-sql",
     # Additional devel dependencies (do not remove this line and add extra development dependencies)
-    "apache-airflow-providers-common-sql[pandas]",
+    "apache-airflow-providers-common-sql[pandas,polars]",
 ]
 
 # To build docs:

providers/sqlite/tests/unit/sqlite/hooks/test_sqlite.py

@@ -99,13 +99,13 @@ def test_get_records(self):
         self.cur.close.assert_called_once_with()
         self.cur.execute.assert_called_once_with(statement)
 
-    def test_get_pandas_df(self):
+    def test_get_df_pandas(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)
+        df = self.db_hook.get_df(statement, df_type="pandas")
 
         assert column == df.columns[0]
 
@@ -114,6 +114,22 @@ def test_get_pandas_df(self):
 
         self.cur.execute.assert_called_once_with(statement)
 
+    def test_get_df_polars(self):
+        statement = "SQL"
+        column = "col"
+        result_sets = [("row1",), ("row2",)]
+        mock_execute = mock.MagicMock()
+        mock_execute.description = [(column, None, None, None, None, None, None)]
+        mock_execute.fetchall.return_value = result_sets
+        self.cur.execute.return_value = mock_execute
+        df = self.db_hook.get_df(statement, df_type="polars")
+
+        self.cur.execute.assert_called_once_with(statement)
+        mock_execute.fetchall.assert_called_once_with()
+        assert column == df.columns[0]
+        assert result_sets[0][0] == df.row(0)[0]
+        assert result_sets[1][0] == df.row(1)[0]
+
     def test_run_log(self):
         statement = "SQL"
         self.db_hook.run(statement)

providers/vertica/README.rst

@@ -54,7 +54,7 @@ Requirements
 PIP package                              Version required
 =======================================  ==================
 ``apache-airflow``                       ``>=2.9.0``
-``apache-airflow-providers-common-sql``  ``>=1.20.0``
+``apache-airflow-providers-common-sql``  ``>=1.26.0``
 ``vertica-python``                       ``>=0.6.0``
 =======================================  ==================