This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new ec5f19d4999 Add a test verifying that SnowparkOperator closes the
Snowpark (#68775)
ec5f19d4999 is described below
commit ec5f19d4999ecce8a03e7853d9c53b12875729dd
Author: SameerMesiah97 <[email protected]>
AuthorDate: Wed Jul 1 22:12:35 2026 +0100
Add a test verifying that SnowparkOperator closes the Snowpark (#68775)
---
generated/provider_dependencies.json.sha256sum | 1 +
.../unit/snowflake/decorators/test_snowpark.py | 32 ++++++-------
.../unit/snowflake/operators/test_snowpark.py | 52 ++++++++++++++++++----
.../tests/unit/snowflake/utils/test_snowpark.py | 2 +-
4 files changed, 61 insertions(+), 26 deletions(-)
diff --git a/generated/provider_dependencies.json.sha256sum
b/generated/provider_dependencies.json.sha256sum
new file mode 100644
index 00000000000..943fd0fc93e
--- /dev/null
+++ b/generated/provider_dependencies.json.sha256sum
@@ -0,0 +1 @@
+93831555f2a141e481c81c147142aeb860c34ea860163ca130d045e5ecd0a83b
diff --git
a/providers/snowflake/tests/unit/snowflake/decorators/test_snowpark.py
b/providers/snowflake/tests/unit/snowflake/decorators/test_snowpark.py
index f0b36e7b41f..b3af126763f 100644
--- a/providers/snowflake/tests/unit/snowflake/decorators/test_snowpark.py
+++ b/providers/snowflake/tests/unit/snowflake/decorators/test_snowpark.py
@@ -23,7 +23,7 @@ from unittest import mock
import pytest
-pytest.importorskip("snowflake-snowpark-python")
+pytest.importorskip("snowflake.snowpark")
from airflow.providers.common.compat.sdk import task
@@ -72,12 +72,12 @@ class TestSnowparkDecorator:
def func2():
return number
- with dag_maker(dag_id=TEST_DAG_ID):
+ with dag_maker(dag_id=TEST_DAG_ID) as dag:
_ = [func1(), func2()]
dr = dag_maker.create_dagrun()
- for ti in dr.get_task_instances():
- ti.run()
+ for task_obj in dag.tasks:
+ ti = dag_maker.run_ti(task_obj.task_id, dr)
assert ti.xcom_pull() == number
assert mock_snowflake_hook.call_count == 2
assert
mock_snowflake_hook.return_value.get_snowpark_session.call_count == 2
@@ -124,12 +124,12 @@ class TestSnowparkDecorator:
def func3(number: int):
return number
- with dag_maker(dag_id=TEST_DAG_ID):
+ with dag_maker(dag_id=TEST_DAG_ID) as dag:
_ = [func1(number=number), func2(number=number),
func3(number=number)]
dr = dag_maker.create_dagrun()
- for ti in dr.get_task_instances():
- ti.run()
+ for task_obj in dag.tasks:
+ ti = dag_maker.run_ti(task_obj.task_id, dr)
assert ti.xcom_pull() == number
assert mock_snowflake_hook.call_count == 3
assert
mock_snowflake_hook.return_value.get_snowpark_session.call_count == 3
@@ -148,12 +148,12 @@ class TestSnowparkDecorator:
def func(session: Session):
assert session ==
mock_snowflake_hook.return_value.get_snowpark_session.return_value
- with dag_maker(dag_id=TEST_DAG_ID):
+ with dag_maker(dag_id=TEST_DAG_ID) as dag:
func()
dr = dag_maker.create_dagrun()
- for ti in dr.get_task_instances():
- ti.run()
+ for task_obj in dag.tasks:
+ ti = dag_maker.run_ti(task_obj.task_id, dr)
assert ti.xcom_pull() is None
mock_snowflake_hook.assert_called_once()
mock_snowflake_hook.return_value.get_snowpark_session.assert_called_once()
@@ -182,12 +182,12 @@ class TestSnowparkDecorator:
assert run_task.xcom.get(key="b") == "2"
assert run_task.xcom.get(key="return_value") == {"a": 1, "b": "2"}
else:
- with dag_maker(dag_id=TEST_DAG_ID):
+ with dag_maker(dag_id=TEST_DAG_ID) as dag:
func()
dr = dag_maker.create_dagrun()
- ti = dr.get_task_instances()[0]
- ti.run()
+ task_obj = dag.tasks[0]
+ ti = dag_maker.run_ti(task_obj.task_id, dr)
assert ti.xcom_pull(key="a") == 1
assert ti.xcom_pull(key="b") == "2"
assert ti.xcom_pull() == {"a": 1, "b": "2"}
@@ -217,12 +217,12 @@ class TestSnowparkDecorator:
def func(session: Session):
return session.query_tag
- with dag_maker(dag_id=TEST_DAG_ID):
+ with dag_maker(dag_id=TEST_DAG_ID) as dag:
func()
dr = dag_maker.create_dagrun()
- ti = dr.get_task_instances()[0]
- ti.run()
+ task_obj = dag.tasks[0]
+ ti = dag_maker.run_ti(task_obj.task_id, dr)
query_tag = ti.xcom_pull()
assert query_tag == {
"dag_id": TEST_DAG_ID,
diff --git
a/providers/snowflake/tests/unit/snowflake/operators/test_snowpark.py
b/providers/snowflake/tests/unit/snowflake/operators/test_snowpark.py
index 062b3112b39..38fcdffe13f 100644
--- a/providers/snowflake/tests/unit/snowflake/operators/test_snowpark.py
+++ b/providers/snowflake/tests/unit/snowflake/operators/test_snowpark.py
@@ -23,7 +23,7 @@ from unittest import mock
import pytest
-pytest.importorskip("snowflake-snowpark-python")
+pytest.importorskip("snowflake.snowpark")
from airflow.providers.snowflake.operators.snowpark import SnowparkOperator
from airflow.utils import timezone
@@ -69,9 +69,10 @@ class TestSnowparkOperator:
]
dr = dag_maker.create_dagrun()
- for ti in dr.get_task_instances():
- ti.run()
+ for task in dag.tasks:
+ ti = dag_maker.run_ti(task.task_id, dr)
assert ti.xcom_pull() == number
+
assert mock_snowflake_hook.call_count == 2
assert
mock_snowflake_hook.return_value.get_snowpark_session.call_count == 2
@@ -109,8 +110,8 @@ class TestSnowparkOperator:
]
dr = dag_maker.create_dagrun()
- for ti in dr.get_task_instances():
- ti.run()
+ for task in dag.tasks:
+ ti = dag_maker.run_ti(task.task_id, dr)
assert ti.xcom_pull() == number
assert mock_snowflake_hook.call_count == 3
assert
mock_snowflake_hook.return_value.get_snowpark_session.call_count == 3
@@ -135,8 +136,8 @@ class TestSnowparkOperator:
)
dr = dag_maker.create_dagrun()
- for ti in dr.get_task_instances():
- ti.run()
+ for task in dag.tasks:
+ ti = dag_maker.run_ti(task.task_id, dr)
assert ti.xcom_pull() is None
mock_snowflake_hook.assert_called_once()
mock_snowflake_hook.return_value.get_snowpark_session.assert_called_once()
@@ -170,8 +171,8 @@ class TestSnowparkOperator:
)
dr = dag_maker.create_dagrun()
- ti = dr.get_task_instances()[0]
- ti.run()
+ task = dag.tasks[0]
+ ti = dag_maker.run_ti(task.task_id, dr)
query_tag = ti.xcom_pull()
assert query_tag == {
"dag_id": TEST_DAG_ID,
@@ -179,3 +180,36 @@ class TestSnowparkOperator:
"task_id": TASK_ID,
"operator": "SnowparkOperator",
}
+
+ @mock.patch("airflow.providers.snowflake.operators.snowpark.SnowflakeHook")
+ def test_snowpark_operator_closes_session_on_exception(
+ self,
+ mock_snowflake_hook,
+ dag_maker,
+ ):
+ mock_session =
mock_snowflake_hook.return_value.get_snowpark_session.return_value
+
+ with dag_maker(dag_id=TEST_DAG_ID) as dag:
+
+ def func(session: Session):
+ raise ValueError("boom")
+
+ SnowparkOperator(
+ task_id=TASK_ID,
+ snowflake_conn_id=CONN_ID,
+ python_callable=func,
+ warehouse="test_warehouse",
+ database="test_database",
+ schema="test_schema",
+ role="test_role",
+ authenticator="externalbrowser",
+ dag=dag,
+ )
+
+ dr = dag_maker.create_dagrun()
+ task = dag.tasks[0]
+
+ with pytest.raises(ValueError, match="boom"):
+ dag_maker.run_ti(task.task_id, dr)
+
+ mock_session.close.assert_called_once()
diff --git a/providers/snowflake/tests/unit/snowflake/utils/test_snowpark.py
b/providers/snowflake/tests/unit/snowflake/utils/test_snowpark.py
index 181d8ed8f3a..c8c9c7945ea 100644
--- a/providers/snowflake/tests/unit/snowflake/utils/test_snowpark.py
+++ b/providers/snowflake/tests/unit/snowflake/utils/test_snowpark.py
@@ -18,7 +18,7 @@ from __future__ import annotations
import pytest
-pytest.importorskip("snowflake-snowpark-python")
+pytest.importorskip("snowflake.snowpark")
from airflow.providers.snowflake.utils.snowpark import
inject_session_into_op_kwargs
