kaxil commented on code in PR #70930:
URL: https://github.com/apache/airflow/pull/70930#discussion_r4201496563
##########
providers/snowflake/docs/operators/snowflake_cortex_agent.rst:
##########
@@ -40,16 +36,56 @@ To use this operator, you must do a few things:
Detailed information is available for :doc:`Installation
<apache-airflow:installation/index>`.
- * :doc:`Setup a Snowflake Connection </connections/snowflake>`.
+ * :doc:`Set up a Snowflake connection </connections/snowflake>`.
- * Create a Snowflake Cortex Agent. See the
- `Snowflake Cortex Agents documentation
<https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-agents>`__.
+Use the ``snowflake_conn_id`` argument to specify the connection used by an
+operator. If not specified, ``snowflake_default`` is used.
-Using the Operator
-^^^^^^^^^^^^^^^^^^
+.. note::
+
+ Parameters passed directly to an operator take precedence over corresponding
Review Comment:
None of the four operators reads `database` or `schema` from the connection:
both are required arguments, every endpoint is built from them, and the hook
only takes the host or account plus credentials from the connection. So there's
nothing for these parameters to take precedence over. I'd drop the note, or say
the connection supplies only the account and credentials.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowflake_cortex_agent.py:
##########
@@ -139,3 +139,276 @@ def execute(self, context: Context) -> dict[str, Any]:
tool_resources=self.tool_resources,
timeout=self.timeout,
)
+
+
+class SnowflakeCortexAgentCreateOperator(BaseOperator):
+ """
+ Create a Snowflake Cortex Agent.
+
+ :param database: Database in which to create the Cortex Agent.
+ :param schema: Schema in which to create the Cortex Agent.
+ :param agent_name: Name of the Cortex Agent.
+ :param comment: Optional comment.
+ :param profile: Agent profile configuration. Optional.
+ :param models: Model configuration. Optional.
+ :param instructions: Agent instructions. Optional.
+ :param orchestration: Orchestration configuration. Optional.
+ :param tools: Agent tools. Optional.
+ :param tool_resources: Tool resource configuration. Optional.
+ :param create_mode: Resource creation mode. Accepted values are
+ ``errorIfExists``, ``orReplace`` and ``ifNotExists``.
+ Defaults to ``errorIfExists``.
+ :param timeout: Maximum time in seconds to wait for the request to
+ complete. Defaults to ``600``.
+ :param snowflake_conn_id: Snowflake connection ID. Defaults to
+ ``snowflake_default``.
+ """
+
+ template_fields: Sequence[str] = (
+ "database",
+ "schema",
+ "agent_name",
+ "comment",
+ "profile",
+ "models",
+ "instructions",
+ "orchestration",
+ "tools",
+ "tool_resources",
+ "snowflake_conn_id",
+ )
+
+ ui_color = "#29B5E8"
+
+ def __init__(
+ self,
+ *,
+ database: str,
+ schema: str,
+ agent_name: str,
+ comment: str | None = None,
+ profile: dict[str, Any] | None = None,
+ models: dict[str, Any] | None = None,
+ instructions: dict[str, Any] | None = None,
+ orchestration: dict[str, Any] | None = None,
+ tools: list[dict[str, Any]] | None = None,
+ tool_resources: dict[str, Any] | None = None,
+ create_mode: CreateMode | str = CreateMode.ERROR_IF_EXISTS,
+ timeout: int | None = 600,
+ snowflake_conn_id: str = "snowflake_default",
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+
+ self.database = database
+ self.schema = schema
+ self.agent_name = agent_name
+ self.comment = comment
+ self.profile = profile
+ self.models = models
+ self.instructions = instructions
+ self.orchestration = orchestration
+ self.tools = tools
+ self.tool_resources = tool_resources
+ self.create_mode = CreateMode(create_mode)
+ self.timeout = timeout
+ self.snowflake_conn_id = snowflake_conn_id
+
+ @cached_property
+ def hook(self) -> SnowflakeCortexAgentHook:
+ """Return the Snowflake Cortex Agent hook."""
+ return SnowflakeCortexAgentHook(
+ snowflake_conn_id=self.snowflake_conn_id,
+ )
+
+ def execute(self, context: Context) -> dict[str, Any]:
+ """Create the Snowflake Cortex Agent."""
+ self.log.info(
+ "Creating Snowflake Cortex Agent '%s.%s.%s'.",
+ self.database,
+ self.schema,
+ self.agent_name,
+ )
+
+ return self.hook.create_agent(
+ database=self.database,
+ schema=self.schema,
+ agent_name=self.agent_name,
+ comment=self.comment,
+ profile=self.profile,
+ models=self.models,
+ instructions=self.instructions,
+ orchestration=self.orchestration,
+ tools=self.tools,
+ tool_resources=self.tool_resources,
+ create_mode=self.create_mode,
+ timeout=self.timeout,
+ )
+
+
+class SnowflakeCortexAgentUpdateOperator(BaseOperator):
+ """
+ Update a Snowflake Cortex Agent.
Review Comment:
The hook drops every `None` field from the update payload, so anything left
unset keeps its current value in Snowflake (the hook's `update_agent` docstring
says so). Could this docstring and the update section of the docs page say the
same? Otherwise leaving `tools` unset to clear an agent's tools looks like it
worked: the task succeeds and the agent is unchanged.
##########
providers/snowflake/tests/unit/snowflake/operators/test_snowflake_cortex_agent.py:
##########
@@ -77,30 +82,117 @@ def test_execute(self, mock_run_agent):
assert result == response
- def test_template_fields(self):
- dag = DAG(
- dag_id="test_template_fields",
- start_date=timezone.datetime(2024, 1, 1),
+
+class TestSnowflakeCortexAgentCreateOperator:
+ @mock.patch.object(
+ SnowflakeCortexAgentHook,
+ "create_agent",
+ autospec=True,
+ )
+ def test_execute(self, mock_create_agent):
Review Comment:
With only the required arguments passed, every optional argument reaches the
hook as `None` or its default, so a swapped pass-through like
`models=self.profile` or a hard-coded `if_exists=False` still passes. I worded
the required-only suggestion as a replacement when it should sit next to the
all-arguments case. A `@pytest.mark.parametrize` per operator with a
required-only case and an all-arguments case using distinct values (for example
`create_mode="orReplace"`, `if_exists=True`, `timeout=300`) covers both the
defaults and the pass-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]
