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 3fbf0134fde Remove redundant Snowflake prefixes from Cortex Agent and
SQL API modules (#74300)
3fbf0134fde is described below
commit 3fbf0134fde572b75f074660a6b35bc1790c0e35
Author: SameerMesiah97 <[email protected]>
AuthorDate: Sat Oct 17 19:31:34 2026 +0100
Remove redundant Snowflake prefixes from Cortex Agent and SQL API modules
(#74300)
---
...snowflake_cortex_agent.rst => cortex_agent.rst} | 8 +++----
.../operators/__init__.py => docs/redirects.txt} | 2 ++
providers/snowflake/provider.yaml | 10 ++++----
.../providers/snowflake/get_provider_info.py | 10 ++++----
.../airflow/providers/snowflake/hooks/__init__.py | 19 +++++++++++++++
.../{snowflake_cortex_agent.py => cortex_agent.py} | 0
.../hooks/{snowflake_sql_api.py => sql_api.py} | 0
.../providers/snowflake/operators/__init__.py | 14 +++++++++++
.../{snowflake_cortex_agent.py => cortex_agent.py} | 2 +-
.../providers/snowflake/operators/snowflake.py | 2 +-
.../snowflake/triggers/snowflake_trigger.py | 2 +-
.../providers/snowflake/utils/openlineage.py | 4 ++--
.../snowflake/example_snowflake_cortex_agent.py | 4 ++--
...wflake_cortex_agent.py => test_cortex_agent.py} | 27 ++++++++++++++++++++--
.../{test_snowflake_sql_api.py => test_sql_api.py} | 4 ++--
...wflake_cortex_agent.py => test_cortex_agent.py} | 4 ++--
.../unit/snowflake/operators/test_snowflake.py | 14 +++++------
.../{test_snowflake_sql.py => test_sql.py} | 0
.../unit/snowflake/triggers/test_snowflake.py | 8 +++----
.../tests/unit/snowflake/utils/test_openlineage.py | 14 +++++------
20 files changed, 103 insertions(+), 45 deletions(-)
diff --git a/providers/snowflake/docs/operators/snowflake_cortex_agent.rst
b/providers/snowflake/docs/operators/cortex_agent.rst
similarity index 91%
rename from providers/snowflake/docs/operators/snowflake_cortex_agent.rst
rename to providers/snowflake/docs/operators/cortex_agent.rst
index 4f7d182addf..8ab07d2971f 100644
--- a/providers/snowflake/docs/operators/snowflake_cortex_agent.rst
+++ b/providers/snowflake/docs/operators/cortex_agent.rst
@@ -66,7 +66,7 @@ SnowflakeCortexAgentCreateOperator
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
To create a Snowflake Cortex Agent you can use
-:class:`~airflow.providers.snowflake.operators.snowflake_cortex_agent.SnowflakeCortexAgentCreateOperator`.
+:class:`~airflow.providers.snowflake.operators.cortex_agent.SnowflakeCortexAgentCreateOperator`.
.. exampleinclude::
/../../snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py
:language: python
@@ -80,7 +80,7 @@ SnowflakeCortexAgentUpdateOperator
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
To update an existing Snowflake Cortex Agent you can use
-:class:`~airflow.providers.snowflake.operators.snowflake_cortex_agent.SnowflakeCortexAgentUpdateOperator`.
+:class:`~airflow.providers.snowflake.operators.cortex_agent.SnowflakeCortexAgentUpdateOperator`.
Only fields explicitly provided are updated. Optional fields left as ``None``
retain their existing values on the Cortex Agent.
@@ -96,7 +96,7 @@ retain their existing values on the Cortex Agent.
SnowflakeCortexAgentOperator
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-Use the
:class:`~airflow.providers.snowflake.operators.snowflake_cortex_agent.SnowflakeCortexAgentOperator`
+Use the
:class:`~airflow.providers.snowflake.operators.cortex_agent.SnowflakeCortexAgentOperator`
to execute an existing Snowflake Cortex Agent.
The operator wraps the Snowflake Cortex Agent Run API and executes an existing
@@ -118,7 +118,7 @@ SnowflakeCortexAgentDeleteOperator
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
To delete a Snowflake Cortex Agent you can use
-:class:`~airflow.providers.snowflake.operators.snowflake_cortex_agent.SnowflakeCortexAgentDeleteOperator`.
+:class:`~airflow.providers.snowflake.operators.cortex_agent.SnowflakeCortexAgentDeleteOperator`.
.. exampleinclude::
/../../snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py
:language: python
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/operators/__init__.py
b/providers/snowflake/docs/redirects.txt
similarity index 92%
copy from
providers/snowflake/src/airflow/providers/snowflake/operators/__init__.py
copy to providers/snowflake/docs/redirects.txt
index 13a83393a91..bf49744b8dc 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/operators/__init__.py
+++ b/providers/snowflake/docs/redirects.txt
@@ -14,3 +14,5 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+
+operators/snowflake_cortex_agent.rst operators/cortex_agent.rst
diff --git a/providers/snowflake/provider.yaml
b/providers/snowflake/provider.yaml
index 62684ef7590..0f375a3a0f9 100644
--- a/providers/snowflake/provider.yaml
+++ b/providers/snowflake/provider.yaml
@@ -125,7 +125,7 @@ integrations:
- /docs/apache-airflow-providers-snowflake/operators/snowflake.rst
- /docs/apache-airflow-providers-snowflake/operators/snowpark.rst
-
/docs/apache-airflow-providers-snowflake/operators/snowpark_containers.rst
- -
/docs/apache-airflow-providers-snowflake/operators/snowflake_cortex_agent.rst
+ - /docs/apache-airflow-providers-snowflake/operators/cortex_agent.rst
logo: /docs/integration-logos/Snowflake.png
tags: [service]
@@ -135,7 +135,7 @@ operators:
- airflow.providers.snowflake.operators.snowflake
- airflow.providers.snowflake.operators.snowpark
- airflow.providers.snowflake.operators.snowpark_containers
- - airflow.providers.snowflake.operators.snowflake_cortex_agent
+ - airflow.providers.snowflake.operators.cortex_agent
task-decorators:
- class-name: airflow.providers.snowflake.decorators.snowpark.snowpark_task
@@ -158,10 +158,10 @@ dataset-uris:
hooks:
- integration-name: Snowflake
python-modules:
- - airflow.providers.snowflake.hooks.snowflake
- - airflow.providers.snowflake.hooks.snowflake_sql_api
- - airflow.providers.snowflake.hooks.snowflake_cortex_agent
- airflow.providers.snowflake.hooks.cortex_model
+ - airflow.providers.snowflake.hooks.snowflake
+ - airflow.providers.snowflake.hooks.sql_api
+ - airflow.providers.snowflake.hooks.cortex_agent
transfers:
- source-integration-name: Amazon Simple Storage Service (S3)
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py
b/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py
index 5d9c989efa3..e0ed30efd8e 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/get_provider_info.py
@@ -34,7 +34,7 @@ def get_provider_info():
"/docs/apache-airflow-providers-snowflake/operators/snowflake.rst",
"/docs/apache-airflow-providers-snowflake/operators/snowpark.rst",
"/docs/apache-airflow-providers-snowflake/operators/snowpark_containers.rst",
-
"/docs/apache-airflow-providers-snowflake/operators/snowflake_cortex_agent.rst",
+
"/docs/apache-airflow-providers-snowflake/operators/cortex_agent.rst",
],
"logo": "/docs/integration-logos/Snowflake.png",
"tags": ["service"],
@@ -47,7 +47,7 @@ def get_provider_info():
"airflow.providers.snowflake.operators.snowflake",
"airflow.providers.snowflake.operators.snowpark",
"airflow.providers.snowflake.operators.snowpark_containers",
-
"airflow.providers.snowflake.operators.snowflake_cortex_agent",
+ "airflow.providers.snowflake.operators.cortex_agent",
],
}
],
@@ -77,10 +77,10 @@ def get_provider_info():
{
"integration-name": "Snowflake",
"python-modules": [
- "airflow.providers.snowflake.hooks.snowflake",
- "airflow.providers.snowflake.hooks.snowflake_sql_api",
- "airflow.providers.snowflake.hooks.snowflake_cortex_agent",
"airflow.providers.snowflake.hooks.cortex_model",
+ "airflow.providers.snowflake.hooks.snowflake",
+ "airflow.providers.snowflake.hooks.sql_api",
+ "airflow.providers.snowflake.hooks.cortex_agent",
],
}
],
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/hooks/__init__.py
b/providers/snowflake/src/airflow/providers/snowflake/hooks/__init__.py
index 13a83393a91..64c10601b9c 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/hooks/__init__.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/hooks/__init__.py
@@ -14,3 +14,22 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+
+from __future__ import annotations
+
+from airflow.utils.deprecation_tools import add_deprecated_classes
+
+__deprecated_classes = {
+ "snowflake_cortex_agent": {
+ "SnowflakeCortexAgentHook": (
+
"airflow.providers.snowflake.hooks.cortex_agent.SnowflakeCortexAgentHook"
+ ),
+ "CreateMode":
"airflow.providers.snowflake.hooks.cortex_agent.CreateMode",
+ "JsonResponse":
"airflow.providers.snowflake.hooks.cortex_agent.JsonResponse",
+ },
+ "snowflake_sql_api": {
+ "SnowflakeSqlApiHook":
("airflow.providers.snowflake.hooks.sql_api.SnowflakeSqlApiHook"),
+ },
+}
+
+add_deprecated_classes(__deprecated_classes, __name__)
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py
b/providers/snowflake/src/airflow/providers/snowflake/hooks/cortex_agent.py
similarity index 100%
rename from
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py
rename to
providers/snowflake/src/airflow/providers/snowflake/hooks/cortex_agent.py
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py
b/providers/snowflake/src/airflow/providers/snowflake/hooks/sql_api.py
similarity index 100%
rename from
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_sql_api.py
rename to providers/snowflake/src/airflow/providers/snowflake/hooks/sql_api.py
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/operators/__init__.py
b/providers/snowflake/src/airflow/providers/snowflake/operators/__init__.py
index 13a83393a91..08854739cfb 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/operators/__init__.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/operators/__init__.py
@@ -14,3 +14,17 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+
+from __future__ import annotations
+
+from airflow.utils.deprecation_tools import add_deprecated_classes
+
+__deprecated_classes = {
+ "snowflake_cortex_agent": {
+ "SnowflakeCortexAgentOperator": (
+
"airflow.providers.snowflake.operators.cortex_agent.SnowflakeCortexAgentOperator"
+ ),
+ },
+}
+
+add_deprecated_classes(__deprecated_classes, __name__)
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake_cortex_agent.py
b/providers/snowflake/src/airflow/providers/snowflake/operators/cortex_agent.py
similarity index 99%
rename from
providers/snowflake/src/airflow/providers/snowflake/operators/snowflake_cortex_agent.py
rename to
providers/snowflake/src/airflow/providers/snowflake/operators/cortex_agent.py
index f631bbb79df..22db75f16dc 100644
---
a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake_cortex_agent.py
+++
b/providers/snowflake/src/airflow/providers/snowflake/operators/cortex_agent.py
@@ -23,7 +23,7 @@ from functools import cached_property
from typing import TYPE_CHECKING, Any
from airflow.providers.common.compat.sdk import BaseOperator
-from airflow.providers.snowflake.hooks.snowflake_cortex_agent import
CreateMode, SnowflakeCortexAgentHook
+from airflow.providers.snowflake.hooks.cortex_agent import CreateMode,
SnowflakeCortexAgentHook
if TYPE_CHECKING:
from airflow.providers.common.compat.sdk import Context
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py
b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py
index dc5aa5c9711..9cb55b1aed9 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py
@@ -33,7 +33,7 @@ from airflow.providers.common.sql.operators.sql import (
SQLIntervalCheckOperator,
SQLValueCheckOperator,
)
-from airflow.providers.snowflake.hooks.snowflake_sql_api import
SnowflakeSqlApiHook
+from airflow.providers.snowflake.hooks.sql_api import SnowflakeSqlApiHook
from airflow.providers.snowflake.triggers.snowflake_trigger import
SnowflakeSqlApiTrigger
_DURABLE_UNSET = object()
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/triggers/snowflake_trigger.py
b/providers/snowflake/src/airflow/providers/snowflake/triggers/snowflake_trigger.py
index 268fb7ac085..b35df38187b 100644
---
a/providers/snowflake/src/airflow/providers/snowflake/triggers/snowflake_trigger.py
+++
b/providers/snowflake/src/airflow/providers/snowflake/triggers/snowflake_trigger.py
@@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Any
from asgiref.sync import sync_to_async
-from airflow.providers.snowflake.hooks.snowflake_sql_api import
SnowflakeSqlApiHook
+from airflow.providers.snowflake.hooks.sql_api import SnowflakeSqlApiHook
from airflow.triggers.base import BaseTrigger, TriggerEvent
if TYPE_CHECKING:
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/utils/openlineage.py
b/providers/snowflake/src/airflow/providers/snowflake/utils/openlineage.py
index f001afec78a..045be13a5bb 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/utils/openlineage.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/utils/openlineage.py
@@ -27,7 +27,7 @@ from airflow.providers.common.compat.sdk import timezone
if TYPE_CHECKING:
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
- from airflow.providers.snowflake.hooks.snowflake_sql_api import
SnowflakeSqlApiHook
+ from airflow.providers.snowflake.hooks.sql_api import SnowflakeSqlApiHook
log = logging.getLogger(__name__)
@@ -168,7 +168,7 @@ def _get_queries_details_from_snowflake(
try:
# Note: need to lazy import here to avoid circular imports
- from airflow.providers.snowflake.hooks.snowflake_sql_api import
SnowflakeSqlApiHook
+ from airflow.providers.snowflake.hooks.sql_api import
SnowflakeSqlApiHook
if isinstance(hook, SnowflakeSqlApiHook):
result = _run_single_query_with_api_hook(hook=hook, sql=query)
diff --git
a/providers/snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py
b/providers/snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py
index 8b579e256cc..16ce0ce0c6e 100644
---
a/providers/snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py
+++
b/providers/snowflake/tests/system/snowflake/example_snowflake_cortex_agent.py
@@ -27,8 +27,8 @@ import os
from datetime import datetime
from airflow import DAG
-from airflow.providers.snowflake.hooks.snowflake_cortex_agent import CreateMode
-from airflow.providers.snowflake.operators.snowflake_cortex_agent import (
+from airflow.providers.snowflake.hooks.cortex_agent import CreateMode
+from airflow.providers.snowflake.operators.cortex_agent import (
SnowflakeCortexAgentCreateOperator,
SnowflakeCortexAgentDeleteOperator,
SnowflakeCortexAgentOperator,
diff --git
a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_cortex_agent.py
b/providers/snowflake/tests/unit/snowflake/hooks/test_cortex_agent.py
similarity index 97%
rename from
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_cortex_agent.py
rename to providers/snowflake/tests/unit/snowflake/hooks/test_cortex_agent.py
index 2c44926506b..5451ef9ab47 100644
---
a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_cortex_agent.py
+++ b/providers/snowflake/tests/unit/snowflake/hooks/test_cortex_agent.py
@@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations
+from importlib import import_module
from unittest import mock
import pytest
@@ -23,13 +24,20 @@ import requests
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.asymmetric import rsa
-from airflow.providers.snowflake.hooks.snowflake_cortex_agent import (
+import airflow.utils.deprecation_tools as deprecation_tools
+from airflow.providers.snowflake.hooks.cortex_agent import (
CreateMode,
JsonResponse,
SnowflakeCortexAgentHook,
)
-MODULE_PATH = "airflow.providers.snowflake.hooks.snowflake_cortex_agent"
+deprecated_import_warning: type[Warning] = getattr(
+ deprecation_tools,
+ "DeprecatedImportWarning",
+ DeprecationWarning,
+)
+
+MODULE_PATH = "airflow.providers.snowflake.hooks.cortex_agent"
HOOK_PATH = f"{MODULE_PATH}.SnowflakeCortexAgentHook"
ACCOUNT = "test-account"
@@ -918,3 +926,18 @@ class TestSnowflakeCortexAgentHook:
with pytest.raises(ValueError, match="Invalid Snowflake account"):
hook._get_base_url()
+
+
[email protected](
+ ("name", "expected"),
+ [
+ ("SnowflakeCortexAgentHook", SnowflakeCortexAgentHook),
+ ("CreateMode", CreateMode),
+ ("JsonResponse", JsonResponse),
+ ],
+)
+def test_deprecated_import(name, expected):
+ old_module =
import_module("airflow.providers.snowflake.hooks.snowflake_cortex_agent")
+
+ with pytest.warns(deprecated_import_warning, match=name):
+ assert getattr(old_module, name) is expected
diff --git
a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_sql_api.py
b/providers/snowflake/tests/unit/snowflake/hooks/test_sql_api.py
similarity index 99%
rename from
providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_sql_api.py
rename to providers/snowflake/tests/unit/snowflake/hooks/test_sql_api.py
index e1a31be539b..0bd6085cb46 100644
--- a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake_sql_api.py
+++ b/providers/snowflake/tests/unit/snowflake/hooks/test_sql_api.py
@@ -34,7 +34,7 @@ from cryptography.hazmat.primitives.asymmetric import rsa
from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.models import Connection
-from airflow.providers.snowflake.hooks.snowflake_sql_api import
SnowflakeSqlApiHook
+from airflow.providers.snowflake.hooks.sql_api import SnowflakeSqlApiHook
if TYPE_CHECKING:
from pathlib import Path
@@ -173,7 +173,7 @@ HOOK_PARAMS: dict = {
API_URL = "https://test.snowflakecomputing.com/api/v2/statements/test";
-MODULE_PATH = "airflow.providers.snowflake.hooks.snowflake_sql_api"
+MODULE_PATH = "airflow.providers.snowflake.hooks.sql_api"
HOOK_PATH = f"{MODULE_PATH}.SnowflakeSqlApiHook"
diff --git
a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake_cortex_agent.py
b/providers/snowflake/tests/unit/snowflake/operators/test_cortex_agent.py
similarity index 98%
rename from
providers/snowflake/tests/unit/snowflake/operators/test_snowflake_cortex_agent.py
rename to
providers/snowflake/tests/unit/snowflake/operators/test_cortex_agent.py
index 226e1a224e0..6ae55a1d7fd 100644
---
a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake_cortex_agent.py
+++ b/providers/snowflake/tests/unit/snowflake/operators/test_cortex_agent.py
@@ -21,11 +21,11 @@ from unittest import mock
import pytest
-from airflow.providers.snowflake.hooks.snowflake_cortex_agent import (
+from airflow.providers.snowflake.hooks.cortex_agent import (
CreateMode,
SnowflakeCortexAgentHook,
)
-from airflow.providers.snowflake.operators.snowflake_cortex_agent import (
+from airflow.providers.snowflake.operators.cortex_agent import (
SnowflakeCortexAgentCreateOperator,
SnowflakeCortexAgentDeleteOperator,
SnowflakeCortexAgentOperator,
diff --git
a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py
b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py
index c8a75157fe8..6f59f7a6b2d 100644
--- a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py
+++ b/providers/snowflake/tests/unit/snowflake/operators/test_snowflake.py
@@ -58,7 +58,7 @@ CONN_ID = "my_snowflake_conn"
TEST_SQL = "select * from any;"
NOTEBOOK = "MY_DB.MY_SCHEMA.MY_NOTEBOOK"
-HOOK_MODULE =
"airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook"
+HOOK_MODULE = "airflow.providers.snowflake.hooks.sql_api.SnowflakeSqlApiHook"
SQL_MULTIPLE_STMTS = (
"create or replace table user_test (i int); insert into user_test (i) "
@@ -438,7 +438,7 @@ class TestSnowflakeSqlApiOperator:
("mock_sql", "statement_count"),
[pytest.param(SQL_MULTIPLE_STMTS, 4, id="multi"),
pytest.param(SINGLE_STMT, 1, id="single")],
)
-
@mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.execute_query")
+
@mock.patch("airflow.providers.snowflake.hooks.sql_api.SnowflakeSqlApiHook.execute_query")
def test_snowflake_sql_api_execute_operator_async(
self, mock_execute_query, mock_sql, statement_count,
mock_get_sql_api_query_status
):
@@ -520,7 +520,7 @@ class TestSnowflakeSqlApiOperator:
({"status": "success", "statement_query_ids": ["uuid", "uuid"]}),
],
)
-
@mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.check_query_output")
+
@mock.patch("airflow.providers.snowflake.hooks.sql_api.SnowflakeSqlApiHook.check_query_output")
def test_snowflake_sql_api_execute_complete(self, mock_conn, mock_event):
"""Tests execute_complete assert with successful message"""
@@ -543,7 +543,7 @@ class TestSnowflakeSqlApiOperator:
({"status": "success", "statement_query_ids": ["uuid", "uuid"]}),
],
)
-
@mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.check_query_output")
+
@mock.patch("airflow.providers.snowflake.hooks.sql_api.SnowflakeSqlApiHook.check_query_output")
def test_snowflake_sql_api_execute_complete_reassigns_query_ids(self,
mock_conn, mock_event):
"""Tests execute_complete assert with successful message"""
@@ -728,7 +728,7 @@ class TestSnowflakeSqlApiOperator:
context["ti"].xcom_push.assert_called_once_with(key="query_ids",
value=["uuid1"])
-
@mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.cancel_queries")
+
@mock.patch("airflow.providers.snowflake.hooks.sql_api.SnowflakeSqlApiHook.cancel_queries")
def test_snowflake_sql_api_on_kill_cancels_queries(self,
mock_cancel_queries):
"""Test that on_kill cancels running queries."""
operator = SnowflakeSqlApiOperator(
@@ -743,7 +743,7 @@ class TestSnowflakeSqlApiOperator:
mock_cancel_queries.assert_called_once_with(["uuid1", "uuid2"])
-
@mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.cancel_queries")
+
@mock.patch("airflow.providers.snowflake.hooks.sql_api.SnowflakeSqlApiHook.cancel_queries")
def test_snowflake_sql_api_on_kill_no_queries(self, mock_cancel_queries):
"""Test that on_kill does nothing when no query ids exist."""
operator = SnowflakeSqlApiOperator(
@@ -758,7 +758,7 @@ class TestSnowflakeSqlApiOperator:
mock_cancel_queries.assert_not_called()
-
@mock.patch("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.cancel_queries")
+
@mock.patch("airflow.providers.snowflake.hooks.sql_api.SnowflakeSqlApiHook.cancel_queries")
def test_snowflake_sql_api_on_kill_respects_cancel_on_kill_false(self,
mock_cancel_queries):
"""on_kill does not cancel queries when cancel_on_kill is disabled."""
operator = SnowflakeSqlApiOperator(
diff --git
a/providers/snowflake/tests/unit/snowflake/operators/test_snowflake_sql.py
b/providers/snowflake/tests/unit/snowflake/operators/test_sql.py
similarity index 100%
rename from
providers/snowflake/tests/unit/snowflake/operators/test_snowflake_sql.py
rename to providers/snowflake/tests/unit/snowflake/operators/test_sql.py
diff --git
a/providers/snowflake/tests/unit/snowflake/triggers/test_snowflake.py
b/providers/snowflake/tests/unit/snowflake/triggers/test_snowflake.py
index 757d5b997ec..2649fe4547e 100644
--- a/providers/snowflake/tests/unit/snowflake/triggers/test_snowflake.py
+++ b/providers/snowflake/tests/unit/snowflake/triggers/test_snowflake.py
@@ -130,7 +130,7 @@ class TestSnowflakeSqlApiTrigger:
@pytest.mark.asyncio
@mock.patch(f"{MODULE}.triggers.snowflake_trigger.SnowflakeSqlApiTrigger.get_query_status")
-
@mock.patch(f"{MODULE}.hooks.snowflake_sql_api.SnowflakeSqlApiHook.get_sql_api_query_status_async")
+
@mock.patch(f"{MODULE}.hooks.sql_api.SnowflakeSqlApiHook.get_sql_api_query_status_async")
async def test_snowflake_sql_trigger_running(
self, mock_get_sql_api_query_status_async, mock_get_query_status
):
@@ -146,7 +146,7 @@ class TestSnowflakeSqlApiTrigger:
@pytest.mark.asyncio
@mock.patch(f"{MODULE}.triggers.snowflake_trigger.SnowflakeSqlApiTrigger.get_query_status")
-
@mock.patch(f"{MODULE}.hooks.snowflake_sql_api.SnowflakeSqlApiHook.get_sql_api_query_status_async")
+
@mock.patch(f"{MODULE}.hooks.sql_api.SnowflakeSqlApiHook.get_sql_api_query_status_async")
async def test_snowflake_sql_trigger_completed(
self, mock_get_sql_api_query_status_async, mock_get_query_status
):
@@ -167,7 +167,7 @@ class TestSnowflakeSqlApiTrigger:
assert TriggerEvent({"status": "success", "statement_query_ids":
QUERY_IDS}) == actual
@pytest.mark.asyncio
-
@mock.patch(f"{MODULE}.hooks.snowflake_sql_api.SnowflakeSqlApiHook.get_sql_api_query_status_async")
+
@mock.patch(f"{MODULE}.hooks.sql_api.SnowflakeSqlApiHook.get_sql_api_query_status_async")
async def test_snowflake_sql_trigger_failure_status(self,
mock_get_sql_api_query_status_async):
"""Test SnowflakeSqlApiTrigger task is executed and triggered with
failure status."""
mock_response = {
@@ -182,7 +182,7 @@ class TestSnowflakeSqlApiTrigger:
assert TriggerEvent(mock_response) == actual
@pytest.mark.asyncio
-
@mock.patch(f"{MODULE}.hooks.snowflake_sql_api.SnowflakeSqlApiHook.get_sql_api_query_status_async")
+
@mock.patch(f"{MODULE}.hooks.sql_api.SnowflakeSqlApiHook.get_sql_api_query_status_async")
async def test_snowflake_sql_trigger_exception(self,
mock_get_sql_api_query_status_async):
"""Tests the SnowflakeSqlApiTrigger does not fire if there is an
exception."""
mock_get_sql_api_query_status_async.side_effect = Exception("Test
exception")
diff --git a/providers/snowflake/tests/unit/snowflake/utils/test_openlineage.py
b/providers/snowflake/tests/unit/snowflake/utils/test_openlineage.py
index 06a42d3abc9..0954eb068a7 100644
--- a/providers/snowflake/tests/unit/snowflake/utils/test_openlineage.py
+++ b/providers/snowflake/tests/unit/snowflake/utils/test_openlineage.py
@@ -26,7 +26,7 @@ import time_machine
from airflow.providers.common.compat.openlineage.facet import SQLJobFacet
from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, timezone
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
-from airflow.providers.snowflake.hooks.snowflake_sql_api import
SnowflakeSqlApiHook
+from airflow.providers.snowflake.hooks.sql_api import SnowflakeSqlApiHook
from airflow.providers.snowflake.utils.openlineage import (
_get_queries_details_from_snowflake,
_process_data_from_api,
@@ -173,10 +173,10 @@ def test_run_single_query_with_hook(mock_get_cursor,
mock_set_autocommit, mock_g
@mock.patch(
-
"airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.get_result_from_successful_sql_api_query"
+
"airflow.providers.snowflake.hooks.sql_api.SnowflakeSqlApiHook.get_result_from_successful_sql_api_query"
)
[email protected]("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.wait_for_query")
[email protected]("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.execute_query")
[email protected]("airflow.providers.snowflake.hooks.sql_api.SnowflakeSqlApiHook.wait_for_query")
[email protected]("airflow.providers.snowflake.hooks.sql_api.SnowflakeSqlApiHook.execute_query")
def test_run_single_query_with_api_hook_success(mock_execute, mock_wait,
mock_get_result):
hook = SnowflakeSqlApiHook(snowflake_conn_id="test_conn")
hook.query_ids = ["old-id"]
@@ -199,10 +199,10 @@ def
test_run_single_query_with_api_hook_success(mock_execute, mock_wait, mock_ge
@mock.patch(
-
"airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.get_result_from_successful_sql_api_query"
+
"airflow.providers.snowflake.hooks.sql_api.SnowflakeSqlApiHook.get_result_from_successful_sql_api_query"
)
[email protected]("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.wait_for_query")
[email protected]("airflow.providers.snowflake.hooks.snowflake_sql_api.SnowflakeSqlApiHook.execute_query")
[email protected]("airflow.providers.snowflake.hooks.sql_api.SnowflakeSqlApiHook.wait_for_query")
[email protected]("airflow.providers.snowflake.hooks.sql_api.SnowflakeSqlApiHook.execute_query")
def test_run_single_query_exception_restores_query_ids(mock_execute,
mock_wait, mock_get_result):
hook = SnowflakeSqlApiHook(snowflake_conn_id="test_conn")
hook.query_ids = ["persistent-id"]
