kaxil commented on code in PR #70101:
URL: https://github.com/apache/airflow/pull/70101#discussion_r3982135395
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py:
##########
@@ -77,7 +108,22 @@ def _request(
response.raise_for_status()
- return response.json()
+ if not response.content:
+ if response_type == "dict":
+ return {}
Review Comment:
Does this need to handle empty bodies at all? None of the three endpoints
return one: per the Snowflake docs `DELETE .../agents/{name}` returns
`{"status": "Request successfully completed"}`, describe returns the agent
object, and `:run` with `stream: false` returns a single populated object.
Where it does fire is `run_agent`, which used to raise `JSONDecodeError` on an
empty 200 and now returns `{}` that `SnowflakeCortexAgentOperator` pushes to
XCom as a successful task with no agent output. Since it keys on
`response_type` rather than on the method, dropping it keeps that failure loud.
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py:
##########
@@ -164,6 +209,111 @@ def run_agent(
endpoint=endpoint,
payload=payload,
timeout=timeout,
+ response_type="dict",
+ )
+
+ def describe_agent(
+ self,
+ *,
+ database: str,
+ schema: str,
+ agent_name: str,
+ timeout: int | None = 600,
+ ) -> JsonDict:
+ """
+ Describe a Snowflake Cortex Agent.
+
+ :param database: Database containing the Cortex Agent.
+ :param schema: Schema containing the Cortex Agent.
+ :param agent_name: Name of the Cortex Agent.
+ :param timeout: Maximum time in seconds to wait for the Cortex Agent
+ request to complete. Optional. Defaults to ``600``.
+ :return: JSON description of the Cortex Agent.
+ """
+ endpoint =
f"/api/v2/databases/{database}/schemas/{schema}/agents/{agent_name}"
+
+ return self._request(
+ method="GET",
+ endpoint=endpoint,
+ timeout=timeout,
+ response_type="dict",
+ )
+
+ def list_agents(
+ self,
+ *,
+ database: str,
+ schema: str,
+ like: str | None = None,
+ from_name: str | None = None,
+ show_limit: int | None = None,
+ timeout: int | None = 600,
+ ) -> JsonList:
+ """
+ List Snowflake Cortex Agents.
+
+ :param database: Database containing the Cortex Agents.
+ :param schema: Schema containing the Cortex Agents.
+ :param like: Case-insensitive name filter. Optional.
+ Defaults to ``None``.
+ :param from_name: Pagination starting point. Optional.
+ Defaults to ``None``.
+ :param show_limit: Maximum number of agents to return. Optional.
+ Defaults to ``None``.
+ :param timeout: Maximum time in seconds to wait for the Cortex Agent
+ request to complete. Optional. Defaults to ``600``.
+ :return: List of Cortex Agents.
+ """
+ endpoint = f"/api/v2/databases/{database}/schemas/{schema}/agents"
+
+ params: dict[str, Any] = {}
+
+ if like is not None:
+ params["like"] = like
+
+ if from_name is not None:
+ params["fromName"] = from_name
+
+ if show_limit is not None:
+ params["showLimit"] = show_limit
+
+ return self._request(
+ method="GET",
+ endpoint=endpoint,
+ params=params or None,
+ timeout=timeout,
+ response_type="list",
+ )
+
+ def delete_agent(
+ self,
+ *,
+ database: str,
+ schema: str,
+ agent_name: str,
+ if_exists: bool = False,
+ timeout: int | None = 600,
+ ) -> JsonDict:
+ """
+ Delete a Snowflake Cortex Agent.
+
+ :param database: Database containing the Cortex Agent.
+ :param schema: Schema containing the Cortex Agent.
+ :param agent_name: Name of the Cortex Agent.
+ :param if_exists: If ``True``, do not fail when the agent does not
exist.
+ Optional. Defaults to ``False``.
+ :param timeout: Maximum time in seconds to wait for the Cortex Agent
request
+ to complete. Optional. Defaults to ``600``.
+ :return: JSON response confirming deletion.
+ """
+ endpoint =
f"/api/v2/databases/{database}/schemas/{schema}/agents/{agent_name}"
Review Comment:
These path segments go into the URL unencoded, and `requests` normalises the
result before sending, so the request can land somewhere other than the
arguments describe. With `agent_name="../../../schemas/OTHER/agents/VICTIM"`
the prepared URL collapses to
`/api/v2/databases/MY_DB/schemas/OTHER/agents/VICTIM`, and a name containing
`?` turns the remainder into query params. Snowflake identifiers created
double-quoted can legitimately contain `/`, `?` and `#`, so this is reachable
without anything unusual. Wrapping each segment in `urllib.parse.quote(x,
safe="")` fixes it, and it is worth doing for `describe_agent`, `list_agents`
and the existing `run_agent` endpoint on line 205 in the same pass.
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py:
##########
@@ -77,7 +108,22 @@ def _request(
response.raise_for_status()
- return response.json()
+ if not response.content:
+ if response_type == "dict":
+ return {}
+ raise TypeError("Expected list[dict] response, got empty response")
+
+ data = response.json()
+
+ if response_type == "dict":
+ if not isinstance(data, dict):
+ raise TypeError(f"Expected dict response, got
{type(data).__name__}")
+ return data
+
+ if not isinstance(data, list) or not all(isinstance(item, dict) for
item in data):
+ raise TypeError(f"Expected list[dict] response, got
{type(data).__name__}")
Review Comment:
When the body is a list whose elements are not dicts, this reports `got
list`, which reads as if a list was the wrong shape: `[1, 2, 3]` raises
`Expected list[dict] response, got list`. Splitting the two conditions would
let the element case say what actually failed. That half is also the only
branch in the new validation without a test, since the parametrized case sends
a dict and short-circuits on `not isinstance(data, list)`.
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py:
##########
@@ -164,6 +209,111 @@ def run_agent(
endpoint=endpoint,
payload=payload,
timeout=timeout,
+ response_type="dict",
+ )
+
+ def describe_agent(
+ self,
+ *,
+ database: str,
+ schema: str,
+ agent_name: str,
+ timeout: int | None = 600,
+ ) -> JsonDict:
+ """
+ Describe a Snowflake Cortex Agent.
+
+ :param database: Database containing the Cortex Agent.
+ :param schema: Schema containing the Cortex Agent.
+ :param agent_name: Name of the Cortex Agent.
+ :param timeout: Maximum time in seconds to wait for the Cortex Agent
+ request to complete. Optional. Defaults to ``600``.
+ :return: JSON description of the Cortex Agent.
+ """
+ endpoint =
f"/api/v2/databases/{database}/schemas/{schema}/agents/{agent_name}"
+
+ return self._request(
+ method="GET",
+ endpoint=endpoint,
+ timeout=timeout,
+ response_type="dict",
+ )
+
+ def list_agents(
+ self,
+ *,
+ database: str,
+ schema: str,
+ like: str | None = None,
+ from_name: str | None = None,
+ show_limit: int | None = None,
+ timeout: int | None = 600,
+ ) -> JsonList:
+ """
+ List Snowflake Cortex Agents.
+
+ :param database: Database containing the Cortex Agents.
+ :param schema: Schema containing the Cortex Agents.
+ :param like: Case-insensitive name filter. Optional.
+ Defaults to ``None``.
+ :param from_name: Pagination starting point. Optional.
+ Defaults to ``None``.
+ :param show_limit: Maximum number of agents to return. Optional.
+ Defaults to ``None``.
+ :param timeout: Maximum time in seconds to wait for the Cortex Agent
+ request to complete. Optional. Defaults to ``600``.
+ :return: List of Cortex Agents.
+ """
+ endpoint = f"/api/v2/databases/{database}/schemas/{schema}/agents"
+
+ params: dict[str, Any] = {}
+
+ if like is not None:
+ params["like"] = like
+
+ if from_name is not None:
+ params["fromName"] = from_name
+
+ if show_limit is not None:
+ params["showLimit"] = show_limit
+
+ return self._request(
Review Comment:
Does this endpoint paginate? Snowflake's v2 spec defines a `Link` response
header for list endpoints carrying `next` and `last` rels, and `_request` drops
response headers entirely, so a caller has no way to tell a full result from a
truncated first page. Since `list_agents` is new API, it seems worth deciding
now whether to surface that or to document `from_name` as the way to page
through.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
