This is an automated email from the ASF dual-hosted git repository.
potiuk 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 77609e83c22 Align Snowflake Cortex Agent hook with the Common AI
managed-agent contract (#74672)
77609e83c22 is described below
commit 77609e83c227ce65b0bb0e606ec73d87aca7bc9d
Author: SameerMesiah97 <[email protected]>
AuthorDate: Sun Nov 1 13:56:29 2026 +0000
Align Snowflake Cortex Agent hook with the Common AI managed-agent contract
(#74672)
* Add Common AI managed-agent support to Snowflake Cortex Agent hook
Implement agent resolution, capabilities and stateless invocation.
Test request conversion, validation, error handling and optional imports.
* Import AirflowOptionalProviderFeatureException from common.compat in the
Cortex agent hook
Generated-by: Claude Opus 5
---------
Co-authored-by: Sameer Mesiah <[email protected]>
Co-authored-by: Jarek Potiuk <[email protected]>
---
providers/snowflake/pyproject.toml | 2 +-
.../providers/snowflake/hooks/cortex_agent.py | 182 +++++++-
.../unit/snowflake/hooks/test_cortex_agent.py | 475 ++++++++++++++++++++-
uv.lock | 2 +-
4 files changed, 649 insertions(+), 12 deletions(-)
diff --git a/providers/snowflake/pyproject.toml
b/providers/snowflake/pyproject.toml
index 111b0c26702..43b1b826ade 100644
--- a/providers/snowflake/pyproject.toml
+++ b/providers/snowflake/pyproject.toml
@@ -93,7 +93,7 @@ dependencies = [
# Any change in the dependencies is preserved when the file is regenerated
[project.optional-dependencies]
"common.ai" = [
- "apache-airflow-providers-common-ai[openai]>=0.10.0",
+ "apache-airflow-providers-common-ai[openai]>=0.10.0", # use next version
]
"microsoft.azure" = [
"apache-airflow-providers-microsoft-azure>=12.8.0"
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/hooks/cortex_agent.py
b/providers/snowflake/src/airflow/providers/snowflake/hooks/cortex_agent.py
index 23f056227c7..999ed6d0b23 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/hooks/cortex_agent.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/hooks/cortex_agent.py
@@ -17,20 +17,90 @@
from __future__ import annotations
+import re
from enum import Enum
-from typing import Any, Literal, overload
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, overload
from urllib.parse import quote
import requests
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
from airflow.providers.snowflake.utils._rest_auth import
SnowflakeRestTokenProvider, get_cortex_base_url
+if TYPE_CHECKING:
+ from airflow.providers.common.ai.exceptions import
ManagedAgentInvocationError
+ from airflow.providers.common.ai.managed_agents.base import (
+ BaseManagedAgentHook,
+ ManagedAgentCapabilities,
+ ManagedAgentRef,
+ ManagedAgentRequest,
+ ManagedAgentResponse,
+ )
+else:
+ try:
+ from airflow.providers.common.ai.exceptions import
ManagedAgentInvocationError
+ from airflow.providers.common.ai.managed_agents.base import (
+ BaseManagedAgentHook,
+ ManagedAgentCapabilities,
+ ManagedAgentRef,
+ ManagedAgentResponse,
+ )
+ except ImportError:
+
+ def _needs_common_ai(*args: Any, **kwargs: Any) -> Any:
+ raise AirflowOptionalProviderFeatureException(
+ "Using a Cortex Agent through the managed-agent contract needs
the 'common.ai' "
+ "extra of the snowflake provider: "
+ "pip install 'apache-airflow-providers-snowflake[common.ai]'."
+ )
+
+ class BaseManagedAgentHook:
+ """Stand-in for the optional Common AI contract."""
+
+ agent = _needs_common_ai
+
+ ManagedAgentCapabilities = ManagedAgentRef = ManagedAgentResponse =
_needs_common_ai
+ ManagedAgentInvocationError = _needs_common_ai
+
+
+# Quoted identifiers can contain dots and escaped double quotes. Unquoted
identifiers
+# follow Snowflake's identifier rules and are normalized to uppercase.
+_IDENTIFIER = r'(?:"(?:[^"\x00]|"")+"|[A-Za-z_][A-Za-z0-9_$]*)'
+_AGENT_NAME = re.compile(rf"({_IDENTIFIER})\.({_IDENTIFIER})\.({_IDENTIFIER})")
+
+_RESERVED_OPTIONS = frozenset(
+ {
+ "database",
+ "schema",
+ "agent_name",
+ "agent",
+ "snowflake_conn_id",
+ "connection",
+ "messages",
+ "prompt",
+ "session_id",
+ "timeout",
+ "thread_id",
+ "parent_message_id",
+ "stream",
+ }
+)
+
+# The named-agent run endpoint supports tool selection; configuration changes
belong
+# to update_agent(). Keep the adapter's option surface explicit.
+_SUPPORTED_OPTIONS = frozenset({"tool_choice"})
+_TERMINAL_HTTP_STATUSES = frozenset({400, 401, 403, 404, 405, 422})
+
JsonDict = dict[str, Any]
JsonList = list[JsonDict]
JsonResponse = JsonDict | JsonList
+class _UnexpectedResponseShape(TypeError):
+ """The decoded response does not have the expected JSON structure."""
+
+
class CreateMode(str, Enum):
"""Resource creation modes for Cortex Agents."""
@@ -39,7 +109,7 @@ class CreateMode(str, Enum):
IF_NOT_EXISTS = "ifNotExists"
-class SnowflakeCortexAgentHook(SnowflakeHook):
+class SnowflakeCortexAgentHook(SnowflakeHook, BaseManagedAgentHook):
"""
Hook for interacting with Snowflake Cortex Agents.
@@ -57,12 +127,104 @@ class SnowflakeCortexAgentHook(SnowflakeHook):
The resolved token is cached and renewed the same way as
``SnowflakeSqlApiHook`` -- see
``SnowflakeRestTokenProvider``, one instance per hook, so a long-running
agent reuses the
same JWT within its renewal window.
+
+ With the ``common.ai`` extra installed,
``hook.agent("DATABASE.SCHEMA.NAME")``
+ supports the Common AI managed-agent contract, accepting prompts or message
+ history and returning answer text with the raw response. The only supported
+ vendor option is ``tool_choice``. Agent configuration changes remain
available
+ through ``update_agent``.
"""
+ agent_platform: ClassVar[str] = "snowflake.cortex"
+
def __init__(self, *args: Any, **kwargs: Any) -> None:
super().__init__(*args, **kwargs)
self._rest_token_provider: SnowflakeRestTokenProvider | None = None
+ @staticmethod
+ def _parse_agent_name(agent: str) -> tuple[str, str, str]:
+ """Validate and normalize name components, preserving quoted
identifier syntax."""
+ match = _AGENT_NAME.fullmatch(agent)
+ if match is None:
+ raise ValueError(f"A Cortex agent must be a fully qualified
DATABASE.SCHEMA.NAME, got {agent!r}.")
+
+ database, schema, agent_name = (
+ part if part.startswith('"') else part.upper() for part in
match.groups()
+ )
+ return database, schema, agent_name
+
+ def resolve_agent(self, agent: str) -> ManagedAgentRef:
+ """Return the normalized agent reference without making a network
call."""
+ # Construct first so the optional-dependency error precedes validation.
+ ref = ManagedAgentRef(platform=self.agent_platform, name=agent)
+ parts = self._parse_agent_name(agent)
+ return ManagedAgentRef(platform=ref.platform, name=".".join(parts))
+
+ def get_agent_capabilities(self, agent: str) -> ManagedAgentCapabilities:
+ """Report the stateless adapter's capabilities without making a
network call."""
+ self.resolve_agent(agent)
+ return ManagedAgentCapabilities()
+
+ def invoke_agent(
+ self,
+ agent: str,
+ request: ManagedAgentRequest,
+ ) -> ManagedAgentResponse:
+ """Invoke a named Cortex Agent through the optional Common AI
contract."""
+ # Check the optional dependency before accessing the request or
validating the name.
+ ManagedAgentRef(platform=self.agent_platform, name=agent)
+
+ database, schema, agent_name = self._parse_agent_name(agent)
+
+ if request.session_id is not None:
+ raise ValueError("The Cortex managed-agent adapter does not
support session_id.")
+
+ if reserved := _RESERVED_OPTIONS.intersection(request.vendor_options):
+ raise ValueError(f"vendor_options cannot override contract fields:
{sorted(reserved)}")
+
+ if unsupported := request.vendor_options.keys() - _SUPPORTED_OPTIONS:
+ raise ValueError(f"Unsupported Cortex vendor_options:
{sorted(unsupported)}")
+
+ messages = []
+ for message in request.as_messages():
+ normalized_message = dict(message)
+ if isinstance(normalized_message.get("content"), str):
+ normalized_message["content"] = [{"type": "text", "text":
normalized_message["content"]}]
+ messages.append(normalized_message)
+
+ try:
+ response = self.run_agent(
+ database=database,
+ schema=schema,
+ agent_name=agent_name,
+ messages=messages,
+ timeout=request.timeout,
+ **request.vendor_options,
+ )
+ except requests.HTTPError as exc:
+ if exc.response is not None and exc.response.status_code in
_TERMINAL_HTTP_STATUSES:
+ raise ManagedAgentInvocationError(
+ f"Cortex agent {agent!r} invocation failed with HTTP
{exc.response.status_code}."
+ ) from exc
+ raise
+ except requests.exceptions.JSONDecodeError as exc:
+ raise ManagedAgentInvocationError(
+ f"Cortex agent {agent!r} returned a response that could not be
decoded as JSON."
+ ) from exc
+ except _UnexpectedResponseShape as exc:
+ raise ManagedAgentInvocationError(
+ f"Cortex agent {agent!r} returned an unexpected response
shape: {exc}"
+ ) from exc
+
+ try:
+ text = self.get_text_response(response)
+ except (AttributeError, TypeError) as exc:
+ raise ManagedAgentInvocationError(
+ f"Cortex agent {agent!r} returned invalid response content."
+ ) from exc
+
+ return ManagedAgentResponse(text=text, raw=response)
+
def _get_base_url(self) -> str:
return get_cortex_base_url(self._get_static_conn_params)
@@ -74,7 +236,7 @@ class SnowflakeCortexAgentHook(SnowflakeHook):
endpoint: str,
payload: JsonDict | None = None,
params: JsonDict | None = None,
- timeout: int | None = None,
+ timeout: float | None = None,
response_type: Literal["dict"],
allow_empty: bool = False,
) -> JsonDict: ...
@@ -87,7 +249,7 @@ class SnowflakeCortexAgentHook(SnowflakeHook):
endpoint: str,
payload: JsonDict | None = None,
params: JsonDict | None = None,
- timeout: int | None = None,
+ timeout: float | None = None,
response_type: Literal["list"],
allow_empty: Literal[False] = False,
) -> JsonList: ...
@@ -99,7 +261,7 @@ class SnowflakeCortexAgentHook(SnowflakeHook):
endpoint: str,
payload: JsonDict | None = None,
params: JsonDict | None = None,
- timeout: int | None = None,
+ timeout: float | None = None,
response_type: Literal["dict", "list"],
allow_empty: bool = False,
) -> JsonResponse:
@@ -135,14 +297,16 @@ class SnowflakeCortexAgentHook(SnowflakeHook):
if response_type == "dict":
if not isinstance(data, dict):
- raise TypeError(f"Expected dict response, got
{type(data).__name__}")
+ raise _UnexpectedResponseShape(f"Expected dict response, got
{type(data).__name__}")
return data
if not isinstance(data, list):
- raise TypeError(f"Expected list[dict] response, got
{type(data).__name__}")
+ raise _UnexpectedResponseShape(f"Expected list[dict] response, got
{type(data).__name__}")
if not all(isinstance(item, dict) for item in data):
- raise TypeError("Expected list[dict] response, got list containing
non-dict elements")
+ raise _UnexpectedResponseShape(
+ "Expected list[dict] response, got list containing non-dict
elements"
+ )
return data
@@ -198,7 +362,7 @@ class SnowflakeCortexAgentHook(SnowflakeHook):
orchestration: dict[str, Any] | None = None,
tools: list[dict[str, Any]] | None = None,
tool_resources: dict[str, Any] | None = None,
- timeout: int | None = 600,
+ timeout: float | None = 600,
) -> JsonDict:
"""
Execute a Snowflake Cortex Agent and return the response payload.
diff --git
a/providers/snowflake/tests/unit/snowflake/hooks/test_cortex_agent.py
b/providers/snowflake/tests/unit/snowflake/hooks/test_cortex_agent.py
index 5451ef9ab47..e9d35906ef3 100644
--- a/providers/snowflake/tests/unit/snowflake/hooks/test_cortex_agent.py
+++ b/providers/snowflake/tests/unit/snowflake/hooks/test_cortex_agent.py
@@ -16,7 +16,10 @@
# under the License.
from __future__ import annotations
+import subprocess
+import sys
from importlib import import_module
+from textwrap import dedent
from unittest import mock
import pytest
@@ -37,6 +40,24 @@ deprecated_import_warning: type[Warning] = getattr(
DeprecationWarning,
)
+
+try:
+ from airflow.providers.common.ai.exceptions import
ManagedAgentInvocationError
+ from airflow.providers.common.ai.managed_agents import (
+ ManagedAgentCapabilities,
+ ManagedAgentRef,
+ ManagedAgentRequest,
+ )
+except ImportError:
+ HAS_COMMON_AI = False
+else:
+ HAS_COMMON_AI = True
+
+needs_common_ai = pytest.mark.skipif(
+ not HAS_COMMON_AI,
+ reason="requires apache-airflow-providers-common-ai",
+)
+
MODULE_PATH = "airflow.providers.snowflake.hooks.cortex_agent"
HOOK_PATH = f"{MODULE_PATH}.SnowflakeCortexAgentHook"
@@ -938,6 +959,458 @@ class TestSnowflakeCortexAgentHook:
)
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
+
+
[email protected]
+def managed_agent_request():
+ with (
+ mock.patch(
+ f"{HOOK_PATH}._get_static_conn_params",
+ new_callable=mock.PropertyMock,
+ return_value=STATIC_CONN_PARAMS,
+ ),
+ mock.patch(
+ f"{HOOK_PATH}._get_conn_params",
+ return_value=CONN_PARAMS,
+ ),
+ mock.patch(f"{MODULE_PATH}.requests.request") as mock_request,
+ ):
+ mock_request.return_value = create_response(
+ json_body={"content": [{"type": "text", "text": "Answer"}]},
+ )
+ yield mock_request
+
+
+@needs_common_ai
+class TestSnowflakeCortexAgentManagedAgent:
+ @pytest.mark.parametrize(
+ ("agent", "expected_path"),
+ [
+ pytest.param(
+ "db.schema.agent",
+ "/api/v2/databases/DB/schemas/SCHEMA/agents/AGENT:run",
+ id="unquoted",
+ ),
+ pytest.param(
+ 'db."schema.part"."agent""name"',
+
"/api/v2/databases/DB/schemas/%22schema.part%22/agents/%22agent%22%22name%22:run",
+ id="quoted-dot-and-escaped-quote",
+ ),
+ pytest.param(
+ '"My/Db"."My Schema"."MyAgent"',
+
"/api/v2/databases/%22My%2FDb%22/schemas/%22My%20Schema%22/agents/%22MyAgent%22:run",
+ id="quoted-case-space-and-slash",
+ ),
+ ],
+ )
+ def test_prompt_and_options_reach_the_api(self, managed_agent_request,
agent, expected_path):
+ hook = SnowflakeCortexAgentHook(snowflake_conn_id="mock_conn_id")
+
+ response = hook.agent(agent).invoke(
+ ManagedAgentRequest(
+ prompt="Hello",
+ timeout=1.5,
+ vendor_options={"tool_choice": {"type": "auto"}},
+ )
+ )
+
+ managed_agent_request.assert_called_once_with(
+ method="POST",
+ url=f"https://{ACCOUNT}.snowflakecomputing.com{expected_path}";,
+ headers={
+ "Authorization": f"Bearer {ACCESS_TOKEN}",
+ "X-Snowflake-Authorization-Token-Type": "OAUTH",
+ "Content-Type": "application/json",
+ },
+ json={
+ "messages": [
+ {
+ "role": "user",
+ "content": [{"type": "text", "text": "Hello"}],
+ }
+ ],
+ "stream": False,
+ "tool_choice": {"type": "auto"},
+ },
+ params=None,
+ timeout=1.5,
+ )
+ assert response.text == "Answer"
+ assert response.raw == {
+ "content": [{"type": "text", "text": "Answer"}],
+ }
+
+ @pytest.mark.parametrize(
+ ("agent", "expected_name"),
+ [
+ pytest.param(
+ "db.schema.agent",
+ "DB.SCHEMA.AGENT",
+ id="unquoted",
+ ),
+ pytest.param(
+ '"MyDb"."MySchema"."MyAgent"',
+ '"MyDb"."MySchema"."MyAgent"',
+ id="quoted-case",
+ ),
+ pytest.param(
+ 'db."schema.part"."agent""name"',
+ 'DB."schema.part"."agent""name"',
+ id="quoted-dot-and-escaped-quote",
+ ),
+ pytest.param(
+ 'DB."A.B".C',
+ 'DB."A.B".C',
+ id="dot-in-schema",
+ ),
+ pytest.param(
+ 'DB.A."B.C"',
+ 'DB.A."B.C"',
+ id="dot-in-agent",
+ ),
+ pytest.param(
+ 'DB.SCHEMA."agent"',
+ 'DB.SCHEMA."agent"',
+ id="quoted-lowercase-agent",
+ ),
+ ],
+ )
+ def test_resolve_agent(self, managed_agent_request, agent, expected_name):
+ hook = SnowflakeCortexAgentHook(snowflake_conn_id="mock_conn_id")
+
+ result = hook.resolve_agent(agent)
+
+ assert result == ManagedAgentRef(
+ platform="snowflake.cortex",
+ name=expected_name,
+ )
+ assert hook.resolve_agent(result.name) == result
+ managed_agent_request.assert_not_called()
+
+ @pytest.mark.parametrize(
+ "agent",
+ [
+ pytest.param("", id="empty"),
+ pytest.param("AGENT", id="missing-database-and-schema"),
+ pytest.param("DB.AGENT", id="missing-schema"),
+ pytest.param("DB..AGENT", id="empty-schema"),
+ pytest.param("DB.SCHEMA.AGENT.EXTRA", id="too-many-parts"),
+ pytest.param('DB.SCHEMA."agent', id="unclosed-quote"),
+ pytest.param('DB.SCHEMA.""', id="empty-quoted-name"),
+ ],
+ )
+ def test_resolve_agent_rejects_invalid_names(self, managed_agent_request,
agent):
+ hook = SnowflakeCortexAgentHook(snowflake_conn_id="mock_conn_id")
+
+ with pytest.raises(ValueError, match="fully qualified"):
+ hook.resolve_agent(agent)
+
+ managed_agent_request.assert_not_called()
+
+ def test_get_agent_capabilities(self, managed_agent_request):
+ hook = SnowflakeCortexAgentHook(snowflake_conn_id="mock_conn_id")
+
+ capabilities = hook.get_agent_capabilities("DB.SCHEMA.AGENT")
+
+ assert capabilities == ManagedAgentCapabilities(
+ sessions=False,
+ structured_output=False,
+ usage=False,
+ trace=False,
+ )
+ managed_agent_request.assert_not_called()
+
+ @pytest.mark.parametrize(
+ "string_content",
+ [
+ pytest.param(False, id="text-blocks"),
+ pytest.param(True, id="string-content"),
+ ],
+ )
+ def test_message_history_reaches_the_api(self, managed_agent_request,
string_content):
+ hook = SnowflakeCortexAgentHook(snowflake_conn_id="mock_conn_id")
+ history = [
+ ("user", "Hello"),
+ ("assistant", "Hi"),
+ ("user", "How are you?"),
+ ]
+ messages = [
+ {
+ "role": role,
+ "content": text if string_content else [{"type": "text",
"text": text}],
+ }
+ for role, text in history
+ ]
+ expected_messages = [
+ {
+ "role": role,
+ "content": [{"type": "text", "text": text}],
+ }
+ for role, text in history
+ ]
+
+ response =
hook.agent("DB.SCHEMA.AGENT").invoke(ManagedAgentRequest(messages=messages))
+
+ assert managed_agent_request.call_args.kwargs["json"] == {
+ "messages": expected_messages,
+ "stream": False,
+ }
+ assert response.text == "Answer"
+
+ # Converting string content must not mutate the caller's messages.
+ assert messages == [
+ {
+ "role": role,
+ "content": text if string_content else [{"type": "text",
"text": text}],
+ }
+ for role, text in history
+ ]
+
+ @pytest.mark.parametrize(
+ "bound",
+ [
+ pytest.param(False, id="direct"),
+ pytest.param(True, id="bound"),
+ ],
+ )
+ def test_session_id_is_rejected_before_the_call(self,
managed_agent_request, bound):
+ hook = SnowflakeCortexAgentHook(snowflake_conn_id="mock_conn_id")
+ request = ManagedAgentRequest(prompt="Hello", session_id="123")
+
+ invoke = (
+ hook.agent("DB.SCHEMA.AGENT").invoke
+ if bound
+ else lambda request: hook.invoke_agent("DB.SCHEMA.AGENT", request)
+ )
+
+ with pytest.raises(ValueError, match="session_id"):
+ invoke(request)
+
+ managed_agent_request.assert_not_called()
+
+ @pytest.mark.parametrize(
+ ("options", "expected_error"),
+ [
+ pytest.param(
+ {"database": "OTHER_DB"},
+ "cannot override contract fields",
+ id="database",
+ ),
+ pytest.param(
+ {"schema": "OTHER_SCHEMA"},
+ "cannot override contract fields",
+ id="schema",
+ ),
+ pytest.param(
+ {"agent_name": "OTHER_AGENT"},
+ "cannot override contract fields",
+ id="agent-name",
+ ),
+ pytest.param(
+ {"snowflake_conn_id": "other_conn"},
+ "cannot override contract fields",
+ id="connection",
+ ),
+ pytest.param(
+ {"messages": []},
+ "cannot override contract fields",
+ id="messages",
+ ),
+ pytest.param(
+ {"timeout": 30},
+ "cannot override contract fields",
+ id="timeout",
+ ),
+ pytest.param(
+ {"thread_id": 123},
+ "cannot override contract fields",
+ id="thread-id",
+ ),
+ pytest.param(
+ {"parent_message_id": 456},
+ "cannot override contract fields",
+ id="parent-message-id",
+ ),
+ pytest.param(
+ {"stream": True},
+ "cannot override contract fields",
+ id="stream",
+ ),
+ pytest.param(
+ {"unknown_option": True},
+ "Unsupported Cortex vendor_options",
+ id="unknown-option",
+ ),
+ ],
+ )
+ def test_vendor_options_are_rejected_before_the_call(
+ self, managed_agent_request, options, expected_error
+ ):
+ hook = SnowflakeCortexAgentHook(snowflake_conn_id="mock_conn_id")
+
+ with pytest.raises(ValueError, match=expected_error):
+
hook.agent("DB.SCHEMA.AGENT").invoke(ManagedAgentRequest(prompt="Hello",
vendor_options=options))
+
+ managed_agent_request.assert_not_called()
+
+ @pytest.mark.parametrize(
+ ("status_code", "terminal"),
+ [
+ pytest.param(400, True, id="bad-request"),
+ pytest.param(401, True, id="unauthorized"),
+ pytest.param(403, True, id="forbidden"),
+ pytest.param(404, True, id="not-found"),
+ pytest.param(405, True, id="method-not-allowed"),
+ pytest.param(422, True, id="unprocessable"),
+ pytest.param(408, False, id="request-timeout"),
+ pytest.param(409, False, id="conflict"),
+ pytest.param(429, False, id="throttled"),
+ pytest.param(500, False, id="server-error"),
+ pytest.param(503, False, id="unavailable"),
+ ],
+ )
+ def test_http_errors_are_terminal_or_propagate(self,
managed_agent_request, status_code, terminal):
+ hook = SnowflakeCortexAgentHook(snowflake_conn_id="mock_conn_id")
+ http_response = create_response(status_code=status_code)
+ managed_agent_request.return_value = http_response
+ original_error = http_response.raise_for_status.side_effect
+
+ # Resolve the optional exception inside the test, not during
collection.
+ expected_exception = ManagedAgentInvocationError if terminal else
requests.exceptions.HTTPError
+
+ with pytest.raises(expected_exception) as raised:
+
hook.agent("DB.SCHEMA.AGENT").invoke(ManagedAgentRequest(prompt="Hello"))
+
+ if terminal:
+ assert raised.value.__cause__ is original_error
+ else:
+ assert raised.value is original_error
+
+ managed_agent_request.assert_called_once()
+
+ @pytest.mark.parametrize(
+ "error_type",
+ [
+ requests.exceptions.ConnectionError,
+ requests.exceptions.ReadTimeout,
+ ],
+ )
+ def test_transport_errors_propagate(self, managed_agent_request,
error_type):
+ hook = SnowflakeCortexAgentHook(snowflake_conn_id="mock_conn_id")
+ error = error_type("Request failed")
+ managed_agent_request.side_effect = error
+
+ with pytest.raises(error_type) as raised:
+
hook.agent("DB.SCHEMA.AGENT").invoke(ManagedAgentRequest(prompt="Hello"))
+
+ assert raised.value is error
+ managed_agent_request.assert_called_once()
+
+ def test_invalid_json_is_terminal(self, managed_agent_request):
+ hook = SnowflakeCortexAgentHook(snowflake_conn_id="mock_conn_id")
+ error = requests.exceptions.JSONDecodeError("Invalid JSON", "{", 1)
+ managed_agent_request.return_value.json.side_effect = error
+
+ with pytest.raises(
+ ManagedAgentInvocationError,
+ match="could not be decoded as JSON",
+ ) as raised:
+
hook.agent("DB.SCHEMA.AGENT").invoke(ManagedAgentRequest(prompt="Hello"))
+
+ assert raised.value.__cause__ is error
+
+ @pytest.mark.parametrize(
+ "body",
+ [
+ pytest.param([], id="list"),
+ pytest.param("unexpected", id="string"),
+ ],
+ )
+ def test_invalid_response_shape_is_terminal(self, managed_agent_request,
body):
+ hook = SnowflakeCortexAgentHook(snowflake_conn_id="mock_conn_id")
+ managed_agent_request.return_value = create_response(json_body=body)
+
+ with pytest.raises(
+ ManagedAgentInvocationError,
+ match="unexpected response shape",
+ ) as raised:
+
hook.agent("DB.SCHEMA.AGENT").invoke(ManagedAgentRequest(prompt="Hello"))
+
+ assert isinstance(raised.value.__cause__, TypeError)
+
+ @pytest.mark.parametrize(
+ "body",
+ [
+ pytest.param({"content": None}, id="null-content"),
+ pytest.param({"content": [1]}, id="invalid-block"),
+ pytest.param(
+ {"content": [{"type": "text", "text": None}]},
+ id="invalid-text",
+ ),
+ ],
+ )
+ def test_invalid_response_content_is_terminal(self, managed_agent_request,
body):
+ hook = SnowflakeCortexAgentHook(snowflake_conn_id="mock_conn_id")
+ managed_agent_request.return_value = create_response(json_body=body)
+
+ with pytest.raises(
+ ManagedAgentInvocationError,
+ match="invalid response content",
+ ):
+
hook.agent("DB.SCHEMA.AGENT").invoke(ManagedAgentRequest(prompt="Hello"))
+
+ def test_unrelated_type_error_propagates(self):
+ hook = SnowflakeCortexAgentHook(snowflake_conn_id="mock_conn_id")
+ request = ManagedAgentRequest(prompt="Hello")
+ error = TypeError("Unrelated implementation error")
+
+ with mock.patch.object(hook, "run_agent", autospec=True,
side_effect=error):
+ with pytest.raises(TypeError) as raised:
+ hook.invoke_agent("DB.SCHEMA.AGENT", request)
+
+ assert raised.value is error
+
+
+def test_hook_imports_without_common_ai():
+ script = dedent(
+ """
+ import sys
+ from unittest import mock
+
+ for name in (
+ "airflow.providers.common.ai.exceptions",
+ "airflow.providers.common.ai.managed_agents.base",
+ ):
+ sys.modules[name] = None
+
+ from airflow.providers.snowflake.hooks.cortex_agent import
SnowflakeCortexAgentHook
+ from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+
+ hook = SnowflakeCortexAgentHook(snowflake_conn_id="mock_conn_id")
+
+ for method, args in (
+ ("agent", ("DB.SCHEMA.AGENT",)),
+ ("resolve_agent", ("DB.SCHEMA.AGENT",)),
+ ("get_agent_capabilities", ("DB.SCHEMA.AGENT",)),
+ ("invoke_agent", ("DB.SCHEMA.AGENT", None)),
+ ):
+ try:
+ getattr(hook, method)(*args)
+ except AirflowOptionalProviderFeatureException as exc:
+ assert "snowflake[common.ai]" in str(exc), exc
+ else:
+ raise AssertionError(f"{method} did not raise")
+
+ with mock.patch.object(hook, "_request", return_value={}) as request:
+ assert hook.run_agent(
+ database="DB",
+ schema="SCHEMA",
+ agent_name="AGENT",
+ messages=[],
+ ) == {}
+ request.assert_called_once()
+ """
+ )
+ subprocess.run([sys.executable, "-c", script], check=True)
diff --git a/uv.lock b/uv.lock
index fef622b1fe8..7b33c110852 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1,5 +1,5 @@
version = 1
-revision = 5
+revision = 3
requires-python = ">=3.11"
resolution-markers = [
"python_full_version >= '3.15' and sys_platform == 'win32'",
