Dev-iL commented on code in PR #70101:
URL: https://github.com/apache/airflow/pull/70101#discussion_r3792600397
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py:
##########
@@ -159,11 +161,124 @@ def run_agent(
endpoint =
f"/api/v2/databases/{database}/schemas/{schema}/agents/{agent_name}:run"
- return self._request(
- method="POST",
- endpoint=endpoint,
- payload=payload,
- timeout=timeout,
+ return cast(
+ "dict[str, Any]",
+ self._request(
+ method="POST",
+ endpoint=endpoint,
+ payload=payload,
+ timeout=timeout,
+ ),
+ )
Review Comment:
Not a big fan of `cast`. How do we know for sure that this is not a list of
dicts? Is it an API contract, or can it be guessed from the combination of
input parameters?
I suggest defining overloads for `self._request` (e.g. one configuration of
inputs that returns a dict and one that returns list[dict]), or if that's
impossible, something like
```python
response = self._request(...)
if not isinstance(response, dict):
raise TypeError(f"Expected a dict response, got {
type(response).__name__ }")
return response # guaranteed to be dict now
```
...and similarly for the other methods.
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake_cortex_agent.py:
##########
@@ -17,7 +17,7 @@
from __future__ import annotations
-from typing import Any
+from typing import Any, cast
Review Comment:
Something like this might be a bit cleaner:
```suggestion
from typing import Any, cast
JsonDict = dict[str, Any]
JsonOrList = JsonDict | list[JsonDict]
```
--
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]
