Skip to content

Commit c6ca8a2

Browse files
guan404mingeladkalpotiukWesley Chiu
authored
test: migrate get_pandas_df to get_df in provider test (apache#49339)
* test: migrate get_pandas_df to get_df in provider test Co-authored-by: Elad Kalif <[email protected]> Co-authored-by: Jarek Potiuk <[email protected]> * chore: add polars in deps * fix: install polars in ci --------- Co-authored-by: Elad Kalif <[email protected]> Co-authored-by: Jarek Potiuk <[email protected]> Co-authored-by: Wesley Chiu <[email protected]>
1 parent 0f573ee commit c6ca8a2

File tree

25 files changed

+144
-43
lines changed

25 files changed

+144
-43
lines changed

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,8 @@ python dependencies for the provided package.
106106
+---------------------+-----------------------------------------------------+----------------------------------------------------------------------------+
107107
| pandas | ``pip install 'apache-airflow[pandas]'`` | Install Pandas library compatible with Airflow |
108108
+---------------------+-----------------------------------------------------+----------------------------------------------------------------------------+
109+
| polars | ``pip install 'apache-airflow[polars]'`` | Polars hooks and operators |
110+
+---------------------+-----------------------------------------------------+----------------------------------------------------------------------------+
109111
| rabbitmq | ``pip install 'apache-airflow[rabbitmq]'`` | RabbitMQ support as a Celery backend |
110112
+---------------------+-----------------------------------------------------+----------------------------------------------------------------------------+
111113
| sentry | ``pip install 'apache-airflow[sentry]'`` | Sentry service for application logging and monitoring |

‎airflow-core/pyproject.toml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,7 @@ dependencies = [
139139
# pre-installed providers
140140
"apache-airflow-providers-common-compat>=1.6.0",
141141
"apache-airflow-providers-common-io>=1.5.3",
142-
"apache-airflow-providers-common-sql>=1.25.0",
142+
"apache-airflow-providers-common-sql>=1.26.0",
143143
"apache-airflow-providers-smtp>=2.0.2",
144144
"apache-airflow-providers-standard>=0.4.0",
145145
]
@@ -238,7 +238,7 @@ dev = [
238238
"apache-airflow-task-sdk",
239239
# TODO(potiuk): eventually we do not want any providers nor apache-airflow extras to be needed for
240240
# airflow-core tests
241-
"apache-airflow[pandas]",
241+
"apache-airflow[pandas,polars]",
242242
"apache-airflow-providers-amazon",
243243
"apache-airflow-providers-celery",
244244
"apache-airflow-providers-cncf-kubernetes",

‎providers/apache/drill/README.rst

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ Requirements
5454
PIP package Version required
5555
======================================= ===========================
5656
``apache-airflow`` ``>=2.10.0``
57-
``apache-airflow-providers-common-sql`` ``>=1.20.0``
57+
``apache-airflow-providers-common-sql`` ``>=1.26.0``
5858
``sqlalchemy-drill`` ``>=1.1.0,!=1.1.6,!=1.1.7``
5959
======================================= ===========================
6060

‎providers/apache/drill/pyproject.toml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ requires-python = "~=3.9"
5858
# After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
5959
dependencies = [
6060
"apache-airflow>=2.10.0",
61-
"apache-airflow-providers-common-sql>=1.20.0",
61+
"apache-airflow-providers-common-sql>=1.26.0",
6262
# Workaround until we get https://github.com/JohnOmernik/sqlalchemy-drill/issues/94 fixed.
6363
"sqlalchemy-drill>=1.1.0,!=1.1.6,!=1.1.7",
6464
]
@@ -70,7 +70,7 @@ dev = [
7070
"apache-airflow-devel-common",
7171
"apache-airflow-providers-common-sql",
7272
# Additional devel dependencies (do not remove this line and add extra development dependencies)
73-
"apache-airflow-providers-common-sql[pandas]",
73+
"apache-airflow-providers-common-sql[pandas,polars]",
7474
]
7575

7676
# To build docs:

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

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,17 +90,33 @@ def test_get_records(self):
9090
assert self.cur.close.call_count == 1
9191
self.cur.execute.assert_called_once_with(statement)
9292

93-
def test_get_pandas_df(self):
93+
def test_get_df_pandas(self):
9494
statement = "SQL"
9595
column = "col"
9696
result_sets = [("row1",), ("row2",)]
9797
self.cur.description = [(column,)]
9898
self.cur.fetchall.return_value = result_sets
99-
df = self.db_hook().get_pandas_df(statement)
99+
df = self.db_hook().get_df(statement, df_type="pandas")
100100

101101
assert column == df.columns[0]
102102
for i, item in enumerate(result_sets):
103103
assert item[0] == df.values.tolist()[i][0]
104104
assert self.conn.close.call_count == 1
105105
assert self.cur.close.call_count == 1
106106
self.cur.execute.assert_called_once_with(statement)
107+
108+
def test_get_df_polars(self):
109+
statement = "SQL"
110+
column = "col"
111+
result_sets = [("row1",), ("row2",)]
112+
mock_execute = MagicMock()
113+
mock_execute.description = [(column, None, None, None, None, None, None)]
114+
mock_execute.fetchall.return_value = result_sets
115+
self.cur.execute.return_value = mock_execute
116+
df = self.db_hook().get_df(statement, df_type="polars")
117+
118+
self.cur.execute.assert_called_once_with(statement)
119+
mock_execute.fetchall.assert_called_once_with()
120+
assert column == df.columns[0]
121+
assert result_sets[0][0] == df.row(0)[0]
122+
assert result_sets[1][0] == df.row(1)[0]

‎providers/apache/druid/README.rst

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ Requirements
5454
PIP package Version required
5555
======================================= ==================
5656
``apache-airflow`` ``>=2.10.0``
57-
``apache-airflow-providers-common-sql`` ``>=1.20.0``
57+
``apache-airflow-providers-common-sql`` ``>=1.26.0``
5858
``pydruid`` ``>=0.4.1``
5959
======================================= ==================
6060

‎providers/apache/druid/pyproject.toml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ requires-python = "~=3.9"
5858
# After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
5959
dependencies = [
6060
"apache-airflow>=2.10.0",
61-
"apache-airflow-providers-common-sql>=1.20.0",
61+
"apache-airflow-providers-common-sql>=1.26.0",
6262
"pydruid>=0.4.1",
6363
]
6464

@@ -77,6 +77,7 @@ dev = [
7777
"apache-airflow-providers-apache-hive",
7878
"apache-airflow-providers-common-sql",
7979
# Additional devel dependencies (do not remove this line and add extra development dependencies)
80+
"apache-airflow-providers-common-sql[polars]",
8081
]
8182

8283
# To build docs:

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

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -452,17 +452,31 @@ def test_get_records(self):
452452
assert self.cur.close.call_count == 1
453453
self.cur.execute.assert_called_once_with(statement)
454454

455-
def test_get_pandas_df(self):
455+
def test_get_df_pandas(self):
456456
statement = "SQL"
457457
column = "col"
458458
result_sets = [("row1",), ("row2",)]
459459
self.cur.description = [(column,)]
460460
self.cur.fetchall.return_value = result_sets
461-
df = self.db_hook().get_pandas_df(statement)
461+
df = self.db_hook().get_df(statement, df_type="pandas")
462462

463463
assert column == df.columns[0]
464464
for i, item in enumerate(result_sets):
465465
assert item[0] == df.values.tolist()[i][0]
466466
assert self.conn.close.call_count == 1
467467
assert self.cur.close.call_count == 1
468468
self.cur.execute.assert_called_once_with(statement)
469+
470+
def test_get_df_polars(self):
471+
statement = "SQL"
472+
column = "col"
473+
result_sets = [("row1",), ("row2",)]
474+
mock_execute = MagicMock()
475+
mock_execute.description = [(column, None, None, None, None, None, None)]
476+
mock_execute.fetchall.return_value = result_sets
477+
self.cur.execute.return_value = mock_execute
478+
479+
df = self.db_hook().get_df(statement, df_type="polars")
480+
assert column == df.columns[0]
481+
assert result_sets[0][0] == df.row(0)[0]
482+
assert result_sets[1][0] == df.row(1)[0]

‎providers/apache/impala/README.rst

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ Requirements
5454
PIP package Version required
5555
======================================= ==================
5656
``impyla`` ``>=0.18.0,<1.0``
57-
``apache-airflow-providers-common-sql`` ``>=1.20.0``
57+
``apache-airflow-providers-common-sql`` ``>=1.26.0``
5858
``apache-airflow`` ``>=2.10.0``
5959
======================================= ==================
6060

‎providers/apache/impala/pyproject.toml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ requires-python = "~=3.9"
5858
# After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
5959
dependencies = [
6060
"impyla>=0.18.0,<1.0",
61-
"apache-airflow-providers-common-sql>=1.20.0",
61+
"apache-airflow-providers-common-sql>=1.26.0",
6262
"apache-airflow>=2.10.0",
6363
]
6464

@@ -77,7 +77,7 @@ dev = [
7777
"apache-airflow-providers-common-sql",
7878
# Additional devel dependencies (do not remove this line and add extra development dependencies)
7979
"kerberos>=1.3.0",
80-
"apache-airflow-providers-common-sql[pandas]"
80+
"apache-airflow-providers-common-sql[pandas,polars]"
8181
]
8282

8383
# To build docs:

0 commit comments

Comments
 (0)