From ab5f7de9dc265ae380855c9a73fec8853075e09e Mon Sep 17 00:00:00 2001 From: gopidesupavan Date: Wed, 29 Jul 2026 17:37:35 +0100 Subject: [PATCH 1/3] Revert AWS services toolset for common AI provider --- dev/breeze/tests/test_selective_checks.py | 16 +- providers/common/ai/docs/examples.rst | 4 +- providers/common/ai/docs/index.rst | 5 +- providers/common/ai/docs/toolsets.rst | 74 +-- providers/common/ai/pyproject.toml | 9 - .../ai/example_dags/example_aws_toolset.py | 54 -- .../providers/common/ai/toolsets/__init__.py | 10 +- .../providers/common/ai/toolsets/aws.py | 509 ------------------ .../tests/unit/common/ai/toolsets/test_aws.py | 359 ------------ uv.lock | 12 +- 10 files changed, 16 insertions(+), 1036 deletions(-) delete mode 100644 providers/common/ai/src/airflow/providers/common/ai/example_dags/example_aws_toolset.py delete mode 100644 providers/common/ai/src/airflow/providers/common/ai/toolsets/aws.py delete mode 100644 providers/common/ai/tests/unit/common/ai/toolsets/test_aws.py diff --git a/dev/breeze/tests/test_selective_checks.py b/dev/breeze/tests/test_selective_checks.py index 699111954d088..a99d5c7bf863d 100644 --- a/dev/breeze/tests/test_selective_checks.py +++ b/dev/breeze/tests/test_selective_checks.py @@ -1121,7 +1121,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str): pytest.param( ("providers/amazon/src/airflow/providers/amazon/provider.yaml",), { - "selected-providers-list-as-string": "amazon apache.hive cncf.kubernetes common.ai " + "selected-providers-list-as-string": "amazon apache.hive cncf.kubernetes " "common.compat common.messaging common.sql databricks exasol ftp google http imap microsoft.azure " "mongo mysql openlineage postgres salesforce ssh teradata", "all-python-versions": f"['{DEFAULT_PYTHON_MAJOR_MINOR_VERSION}']", @@ -1148,9 +1148,9 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str): { "description": "amazon...google", "test_types": "Providers[amazon] Providers[apache.hive,cncf.kubernetes," - "common.ai,common.compat,common.messaging,common.sql,databricks,exasol,ftp," - "http,imap,microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh," - "teradata] Providers[google]", + "common.compat,common.messaging,common.sql,databricks,exasol,ftp,http,imap," + "microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh,teradata] " + "Providers[google]", } ] ), @@ -1193,7 +1193,7 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str): pytest.param( ("providers/amazon/src/airflow/providers/amazon/file.py",), { - "selected-providers-list-as-string": "amazon apache.hive cncf.kubernetes common.ai " + "selected-providers-list-as-string": "amazon apache.hive cncf.kubernetes " "common.compat common.messaging common.sql databricks exasol ftp google http imap microsoft.azure " "mongo mysql openlineage postgres salesforce ssh teradata", "all-python-versions": f"['{DEFAULT_PYTHON_MAJOR_MINOR_VERSION}']", @@ -1217,9 +1217,9 @@ def assert_outputs_are_printed(expected_outputs: dict[str, str], stderr: str): { "description": "amazon...google", "test_types": "Providers[amazon] Providers[apache.hive,cncf.kubernetes," - "common.ai,common.compat,common.messaging,common.sql,databricks,exasol,ftp," - "http,imap,microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh," - "teradata] Providers[google]", + "common.compat,common.messaging,common.sql,databricks,exasol,ftp,http,imap," + "microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh,teradata] " + "Providers[google]", } ] ), diff --git a/providers/common/ai/docs/examples.rst b/providers/common/ai/docs/examples.rst index 973a4642d2fb7..a74550d5e60b0 100644 --- a/providers/common/ai/docs/examples.rst +++ b/providers/common/ai/docs/examples.rst @@ -69,9 +69,7 @@ Agents & tools `example_agent_capabilities.py `__). * - :ref:`Toolsets ` - Loading ``SKILL.md`` Agent Skills - (`example_agent_skills.py `__), - allow-listed AWS API access via ``AWSToolset`` - (`example_aws_toolset.py `__), + (`example_agent_skills.py `__) and exposing an Airflow toolset to a LangChain agent, the reverse bridge (`example_langchain_toolset_bridge.py `__). * - :doc:`connections/mcp` diff --git a/providers/common/ai/docs/index.rst b/providers/common/ai/docs/index.rst index 7e37a25ef86bd..41a01074f44c9 100644 --- a/providers/common/ai/docs/index.rst +++ b/providers/common/ai/docs/index.rst @@ -217,13 +217,12 @@ You can install such cross-provider dependencies when installing from PyPI. For .. code-block:: bash - pip install apache-airflow-providers-common-ai[amazon] + pip install apache-airflow-providers-common-ai[common.sql] ============================================================================================================ ============== Dependent package Extra ============================================================================================================ ============== -`apache-airflow-providers-amazon `_ ``amazon`` `apache-airflow-providers-common-sql `_ ``common.sql`` `apache-airflow-providers-git `_ ``git`` ============================================================================================================ ============== @@ -253,14 +252,12 @@ Extra Dependencies ``avro`` ``fastavro>=1.10.0; python_version < "3.14"``, ``fastavro>=1.12.1; python_version >= "3.14"`` ``parquet`` ``pyarrow>=18.0.0; python_version < '3.14'``, ``pyarrow>=22.0.0; python_version >= '3.14'`` ``sql`` ``apache-airflow-providers-common-sql>=1.33.0``, ``sqlglot>=30.0.0`` -``aws`` ``apache-airflow-providers-amazon>=9.0.0`` ``common.sql`` ``apache-airflow-providers-common-sql>=1.33.0`` ``langchain`` ``langchain>=1.0.0`` ``llamaindex`` ``dataclasses-json>=0.6.7``, ``llama-index-core>=0.13.0``, ``llama-index-embeddings-openai>=0.6.0``, ``llama-index-llms-openai>=0.6.0`` ``pdf`` ``pypdf>=4.0.0`` ``docx`` ``python-docx>=1.0.0`` ``git`` ``apache-airflow-providers-git`` -``amazon`` ``apache-airflow-providers-amazon`` ============== ======================================================================================================================================= Downloading official packages diff --git a/providers/common/ai/docs/toolsets.rst b/providers/common/ai/docs/toolsets.rst index d14e61772b5a3..a4b08bf869356 100644 --- a/providers/common/ai/docs/toolsets.rst +++ b/providers/common/ai/docs/toolsets.rst @@ -24,20 +24,17 @@ Airflow's 350+ provider hooks already have typed methods, rich docstrings, and managed credentials. Toolsets expose them as pydantic-ai tools so that LLM agents can call them during multi-turn reasoning. -Four toolsets are included: +Three toolsets are included: -- :class:`~airflow.providers.common.ai.toolsets.aws.AWSToolset` — configured - AWS services toolset for agent access to AWS APIs through Airflow-managed - AWS connections. - :class:`~airflow.providers.common.ai.toolsets.hook.HookToolset` — generic adapter for any Airflow Hook. +- :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` — curated + 4-tool database toolset. - :class:`~airflow.providers.common.ai.toolsets.mcp.MCPToolset` — connect to `MCP servers `__ configured via Airflow connections. -- :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` — curated - 4-tool database toolset. -All four implement pydantic-ai's +All three implement pydantic-ai's `AbstractToolset `__ interface and can be passed to any pydantic-ai ``Agent``, including via :class:`~airflow.providers.common.ai.operators.agent.AgentOperator`. @@ -271,69 +268,6 @@ Parameters support DDL for in-memory tables; this guard blocks those by default. - ``max_rows``: Maximum rows returned from the ``query`` tool. Default ``50``. -``AWSToolset`` --------------- - -Curated toolset that gives an agent allow-listed access to AWS APIs with -three tools: - -.. list-table:: - :header-rows: 1 - :widths: 20 50 - - * - Tool - - Description - * - ``list_aws_operations`` - - Lists the operations the toolset allows, grouped by service - * - ``describe_aws_operation`` - - Returns an operation's parameter and response shapes (from the - botocore service model), so the agent can check what a call expects - before making it - * - ``call_aws`` - - Executes an allowed operation and returns the response as JSON, - aggregating paginated results - -Credentials, region, and session configuration come from the Airflow -connection (``aws_conn_id``) via the amazon provider — the model cannot -override them through tool arguments. Requires the ``aws`` extra:: - - pip install "apache-airflow-providers-common-ai[aws]" - -.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_aws_toolset.py - :language: python - :start-after: [START howto_operator_agent_aws] - :end-before: [END howto_operator_agent_aws] - -Access is deny-by-default: ``allowed_actions`` is required and supports ``*`` -wildcards in the operation part only (``"athena:*"``, ``"s3:List*"``). -Operations that return credentials or decrypted secrets — ``sts:AssumeRole``, -``secretsmanager:GetSecretValue``, ``kms:Decrypt``, and similar — are never -matched by a wildcard and must be listed verbatim to be callable. - -.. warning:: - ``allowed_actions`` bounds what the agent can ask for, not what the - credentials can do. Point ``aws_conn_id`` at a least-privilege IAM role - scoped to the same operations. - -Parameters -^^^^^^^^^^ - -- ``aws_conn_id``: Airflow connection ID for AWS credentials. Default - ``aws_default``. -- ``allowed_actions``: Operations the agent may call, in - ``":"`` form using boto3 service names and API - operation names. Required — there is no auto-discovery. Matching is case- - and underscore-insensitive, so ``"s3:list_buckets"`` equals - ``"s3:ListBuckets"``. Validated against botocore's bundled service - definitions at instantiation time (local metadata only — no network, no - credentials). -- ``region_name``: AWS region for API calls. Default ``None`` — use the - region configured on the connection. -- ``max_items``: Upper bound on items aggregated from paginated operations. - Default ``1000``. -- ``max_output_bytes``: Upper bound on the serialized response returned to - the agent; larger payloads are clipped and flagged. Default ``65536``. - ``LoggingToolset`` ------------------ diff --git a/providers/common/ai/pyproject.toml b/providers/common/ai/pyproject.toml index 6254481d08fd7..e97151ecadb55 100644 --- a/providers/common/ai/pyproject.toml +++ b/providers/common/ai/pyproject.toml @@ -108,11 +108,6 @@ dependencies = [ "apache-airflow-providers-common-sql>=1.33.0", "sqlglot>=30.0.0", ] -# AWSToolset: allow-listed AWS API access for agents. The amazon provider -# supplies credential resolution (AwsBaseHook) and brings boto3/botocore. -"aws" = [ - "apache-airflow-providers-amazon>=9.0.0", -] "common.sql" = [ "apache-airflow-providers-common-sql>=1.33.0" ] @@ -130,16 +125,12 @@ dependencies = [ "git" = [ "apache-airflow-providers-git" ] -"amazon" = [ - "apache-airflow-providers-amazon" -] [dependency-groups] dev = [ "apache-airflow", "apache-airflow-task-sdk", "apache-airflow-devel-common", - "apache-airflow-providers-amazon", "apache-airflow-providers-common-compat", "apache-airflow-providers-common-sql", "apache-airflow-providers-git", diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_aws_toolset.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_aws_toolset.py deleted file mode 100644 index 682208447efff..0000000000000 --- a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_aws_toolset.py +++ /dev/null @@ -1,54 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -"""Example Dag: agent with allow-listed AWS access via AWSToolset.""" - -from __future__ import annotations - -from airflow.providers.common.ai.operators.agent import AgentOperator -from airflow.providers.common.ai.toolsets.aws import AWSToolset -from airflow.providers.common.compat.sdk import dag - - -# [START howto_operator_agent_aws] -@dag(tags=["example"]) -def example_agent_aws_toolset(): - AgentOperator( - task_id="s3_auditor", - prompt="Which buckets exist, and roughly how much data is in 'data-lake-raw'?", - llm_conn_id="pydanticai_default", - system_prompt=( - "You are an AWS operations assistant. Discover what you are allowed " - "to call, check parameter shapes before calling, and answer with " - "concrete numbers." - ), - toolsets=[ - AWSToolset( - aws_conn_id="aws_default", - allowed_actions=[ - "s3:ListBuckets", - "s3:ListObjectsV2", - "s3:GetBucketLocation", - ], - region_name="us-east-1", - ) - ], - ) - - -# [END howto_operator_agent_aws] - -example_agent_aws_toolset() diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py index e36690b036df0..6c30fa4a7336c 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py @@ -20,7 +20,7 @@ from airflow.providers.common.ai.toolsets.hook import HookToolset -__all__ = ["AWSToolset", "HookToolset", "MCPToolset", "SQLToolset", "airflow_toolset_to_langchain_tools"] +__all__ = ["HookToolset", "MCPToolset", "SQLToolset", "airflow_toolset_to_langchain_tools"] def __getattr__(name: str): @@ -46,12 +46,4 @@ def __getattr__(name: str): raise AirflowOptionalProviderFeatureException(e) return MCPToolset - if name == "AWSToolset": - try: - from airflow.providers.common.ai.toolsets.aws import AWSToolset - except ImportError as e: - from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException - - raise AirflowOptionalProviderFeatureException() from e - return AWSToolset raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/aws.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/aws.py deleted file mode 100644 index b37adec7497a6..0000000000000 --- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/aws.py +++ /dev/null @@ -1,509 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -"""Curated AWS toolset exposing allow-listed boto3 operations as pydantic-ai tools.""" - -from __future__ import annotations - -import json -import re -from fnmatch import fnmatchcase -from functools import cache -from typing import TYPE_CHECKING, Any - -from pydantic_ai.exceptions import ModelRetry -from pydantic_ai.tools import ToolDefinition -from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool -from pydantic_core import SchemaValidator, core_schema - -from airflow.providers.common.ai.utils.tool_definition import return_schema_kwargs - -try: - import botocore.session - from botocore import xform_name - from botocore.response import StreamingBody - - from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook -except ImportError as e: - from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException - - raise AirflowOptionalProviderFeatureException() from e - -if TYPE_CHECKING: - from botocore.client import BaseClient - from botocore.model import ServiceModel, Shape - from pydantic_ai._run_context import RunContext - -_PASSTHROUGH_VALIDATOR = SchemaValidator(core_schema.any_schema()) - - -def _normalize_action(action: str) -> str: - """Comparison key for actions: case- and underscore-insensitive.""" - return action.replace("_", "").casefold() - - -# High-confidence operations that can put credentials, tokens, passwords, -# decrypted secret material, or plaintext key material into the agent's context. -# A wildcard in ``allowed_actions`` never matches these -- each must be listed -# verbatim to be callable. This is a defense-in-depth guard, not an exhaustive -# AWS security boundary; least-privilege IAM on the connection's role remains -# the real boundary. -_CREDENTIAL_RETURNING_ACTIONS = frozenset( - _normalize_action(action) - for action in ( - "sts:AssumeRole", - "sts:AssumeRoleWithSAML", - "sts:AssumeRoleWithWebIdentity", - "sts:GetFederationToken", - "sts:GetSessionToken", - "secretsmanager:GetSecretValue", - "secretsmanager:BatchGetSecretValue", - "ssm:GetParameter", - "ssm:GetParameters", - "ssm:GetParametersByPath", - "kms:Decrypt", - "kms:GenerateDataKey", - "kms:GenerateDataKeyPair", - "ecr:GetAuthorizationToken", - "ecr-public:GetAuthorizationToken", - "iam:CreateAccessKey", - "iam:CreateLoginProfile", - "iam:UpdateLoginProfile", - "iam:CreateServiceSpecificCredential", - "iam:ResetServiceSpecificCredential", - "cognito-identity:GetCredentialsForIdentity", - "cognito-identity:GetOpenIdToken", - "cognito-identity:GetOpenIdTokenForDeveloperIdentity", - "redshift:GetClusterCredentials", - "redshift:GetClusterCredentialsWithIAM", - "redshift-serverless:GetCredentials", - "lightsail:GetInstanceAccessDetails", - "lightsail:GetRelationalDatabaseMasterUserPassword", - ) -) - -_SHAPE_DEPTH = 3 - -# JSON Schemas for the three AWS tools. -_LIST_OPERATIONS_SCHEMA: dict[str, Any] = { - "type": "object", - "properties": { - "service": { - "type": "string", - "description": "Optional service name filter, e.g. 's3'.", - }, - }, -} - -_DESCRIBE_OPERATION_SCHEMA: dict[str, Any] = { - "type": "object", - "properties": { - "service": {"type": "string", "description": "AWS service name, e.g. 's3'."}, - "operation": { - "type": "string", - "description": "API operation name, e.g. 'ListObjectsV2'.", - }, - }, - "required": ["service", "operation"], -} - -_CALL_AWS_SCHEMA: dict[str, Any] = { - "type": "object", - "properties": { - "service": {"type": "string", "description": "AWS service name, e.g. 's3'."}, - "operation": { - "type": "string", - "description": "API operation name, e.g. 'ListObjectsV2'.", - }, - "parameters": { - "type": "object", - "description": ( - "Operation parameters exactly as the API expects them (see describe_aws_operation)." - ), - }, - }, - "required": ["service", "operation"], -} - - -class AWSToolset(AbstractToolset[Any]): - """ - Curated toolset that gives an LLM agent allow-listed access to AWS APIs. - - Exposes three tools: - - - ``list_aws_operations`` -- the operations this toolset allows, grouped by - service (the allow-list intersected with the botocore service model). - - ``describe_aws_operation`` -- an operation's parameter and response shapes, - derived from the botocore service model, so the agent can check what a - call expects *before* making it. - - ``call_aws`` -- execute an allowed operation and return the response as - JSON (paginated operations are aggregated up to ``max_items``). - - Credentials, region, and session configuration come exclusively from the - Airflow connection (resolved lazily through the amazon provider's - ``AwsBaseHook``); none of them are tool arguments, so the model cannot - steer them. - - When a call fails, AWS's own error message is returned to the agent as a - retry (:class:`pydantic_ai.ModelRetry`) so the model can correct its - parameters within the run. pydantic-ai bounds this by the tool's - ``max_retries``, so an unrecoverable error -- bad credentials, missing IAM - permissions -- exhausts the retries and fails the task for Airflow to retry. - - :param aws_conn_id: Airflow connection ID for AWS credentials. - :param allowed_actions: Operations the agent may call, in - ``":"`` form using boto3 service names and API - operation names -- e.g. ``["s3:ListBuckets", "s3:GetObject", - "athena:*"]``. Required -- access is deny-by-default and there is - deliberately no auto-discovery. The operation part accepts ``*``/``?`` - wildcards; the service part must be explicit. Matching is case- and - underscore-insensitive, so ``"s3:list_buckets"`` and - ``"s3:ListBuckets"`` are equivalent. - - Operations that return credentials or decrypted secrets (for example - ``sts:AssumeRole``, ``secretsmanager:GetSecretValue``, ``kms:Decrypt``) - are never matched by a wildcard -- each must be listed verbatim to be - callable. - - .. note:: - This is an application-level guardrail: it bounds what the agent - can *ask for*, not what the credentials can *do*. Pair it with a - least-privilege IAM role on ``aws_conn_id`` for a hard guarantee. - - :param region_name: AWS region for API calls. ``None`` (default) uses the - region configured on the connection. Deliberately not exposed as a tool - argument. - :param max_items: Upper bound on items aggregated from paginated operations - (applied via the paginator's ``MaxItems``). Default ``1000``. - :param max_output_bytes: Upper bound on the serialized response returned to - the agent. Larger payloads are clipped and flagged with - ``"truncated": true``. Default ``65536``. - """ - - def __init__( - self, - aws_conn_id: str = "aws_default", - *, - allowed_actions: list[str], - region_name: str | None = None, - max_items: int = 1000, - max_output_bytes: int = 65536, - ) -> None: - if not allowed_actions: - raise ValueError("allowed_actions must be a non-empty list.") - - # Validation below uses botocore's bundled service definitions only -- - # local file reads, no credentials, no network -- so it is safe to run - # at Dag parse time. - patterns: dict[str, list[str]] = {} - literal_actions: set[str] = set() - for action in allowed_actions: - service, sep, operation = action.partition(":") - if not sep or not service or not operation: - raise ValueError( - f"Invalid action {action!r}: expected ':', e.g. 's3:ListBuckets'." - ) - if any(ch in service for ch in "*?["): - raise ValueError(f"Invalid action {action!r}: wildcards are not allowed in the service part.") - if service not in _get_available_services(): - raise ValueError( - f"Unknown AWS service {service!r} in action {action!r}. " - "Use the boto3 service name, e.g. 's3' or 'secretsmanager'." - ) - key = _normalize_action(action) - if any(ch in operation for ch in "*?["): - model = _load_service_model(service) - matched_operations = [ - name - for name in model.operation_names - if self._is_pattern_match_allowed(key, service, name) - ] - if not matched_operations: - raise ValueError( - f"Action pattern {action!r} does not match any non-sensitive operation " - f"for service {service!r}." - ) - patterns.setdefault(service, []).append(key) - else: - if _resolve_operation_name(_load_service_model(service), operation) is None: - raise ValueError( - f"Unknown operation {operation!r} for service {service!r} in action " - f"{action!r}. Use the API operation name, e.g. 'ListBuckets'." - ) - patterns.setdefault(service, []).append(key) - literal_actions.add(key) - - self._aws_conn_id = aws_conn_id - self._region_name = region_name - self._max_items = max_items - self._max_output_bytes = max_output_bytes - self._patterns = patterns - self._literal_actions = frozenset(literal_actions) - # Populated lazily at call time, so Dag parsing never creates AWS clients. - self._clients: dict[str, BaseClient] = {} - - @property - def id(self) -> str: - return f"aws-{self._aws_conn_id}" - - # ------------------------------------------------------------------ - # AbstractToolset interface - # ------------------------------------------------------------------ - - async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]: - tools: dict[str, ToolsetTool[Any]] = {} - - for name, description, schema in ( - ( - "list_aws_operations", - "List the AWS operations this toolset allows, grouped by service.", - _LIST_OPERATIONS_SCHEMA, - ), - ( - "describe_aws_operation", - "Show an operation's parameters and response shape before calling it.", - _DESCRIBE_OPERATION_SCHEMA, - ), - ( - "call_aws", - "Execute an allowed AWS operation and return the response as JSON.", - _CALL_AWS_SCHEMA, - ), - ): - # sequential=True because all tools share per-service boto3 clients - # with synchronous I/O -- they must not run concurrently. - # return_schema is "string": every tool returns a JSON-encoded string, - # so code mode renders `-> str` instead of `-> Any`. - tool_def = ToolDefinition( - name=name, - description=description, - parameters_json_schema=schema, - sequential=True, - **return_schema_kwargs({"type": "string"}), - ) - tools[name] = ToolsetTool( - toolset=self, - tool_def=tool_def, - max_retries=1, - args_validator=_PASSTHROUGH_VALIDATOR, - ) - return tools - - async def call_tool( - self, - name: str, - tool_args: dict[str, Any], - ctx: RunContext[Any], - tool: ToolsetTool[Any], - ) -> Any: - if name not in ("list_aws_operations", "describe_aws_operation", "call_aws"): - raise ValueError(f"Unknown tool: {name!r}") - try: - if name == "list_aws_operations": - return self._list_operations(tool_args.get("service")) - if name == "describe_aws_operation": - return self._describe_operation(tool_args["service"], tool_args["operation"]) - return self._call_aws( - tool_args["service"], tool_args["operation"], tool_args.get("parameters") or {} - ) - except Exception as e: - # Hand AWS's own error back to the agent as a retry so it can correct - # its parameters within the run. pydantic-ai bounds this by the tool's - # max_retries, so an unrecoverable error (bad credentials, missing IAM - # permissions) exhausts the budget and fails the task for Airflow to - # retry, rather than being silently worked around. - raise ModelRetry( - f"The {name} tool failed: {e}\n" - "Use list_aws_operations to see what you may call and " - "describe_aws_operation to check the expected parameters, then try again." - ) from e - - # ------------------------------------------------------------------ - # Authorization - # ------------------------------------------------------------------ - - def _is_action_allowed(self, *, service: str, operation: str) -> bool: - key = _normalize_action(f"{service}:{operation}") - if key in _CREDENTIAL_RETURNING_ACTIONS: - return key in self._literal_actions - return any(fnmatchcase(key, pattern) for pattern in self._patterns.get(service, ())) - - def _is_pattern_match_allowed(self, pattern: str, service: str, operation: str) -> bool: - key = _normalize_action(f"{service}:{operation}") - return key not in _CREDENTIAL_RETURNING_ACTIONS and fnmatchcase(key, pattern) - - def _resolve_allowed_service(self, service: str) -> str: - """Validate a tool-supplied service name against the allow-list.""" - if service not in self._patterns: - raise ValueError( - f"Service {service!r} is not in this toolset's allowed actions. " - f"Allowed services: {', '.join(sorted(self._patterns))}." - ) - return service - - def _resolve_allowed_operation(self, service: str, operation: str) -> tuple[str, ServiceModel]: - model = _load_service_model(self._resolve_allowed_service(service)) - canonical = _resolve_operation_name(model, operation) - if canonical is None: - raise ValueError(f"Unknown operation {operation!r} for service {service!r}.") - if not self._is_action_allowed(service=service, operation=canonical): - raise ValueError(f"Operation {service}:{canonical} is not in this toolset's allowed actions.") - return canonical, model - - # ------------------------------------------------------------------ - # Tool implementations - # ------------------------------------------------------------------ - - def _list_operations(self, service: str | None) -> str: - services = [self._resolve_allowed_service(service)] if service else sorted(self._patterns) - listing = { - svc: sorted( - name - for name in _load_service_model(svc).operation_names - if self._is_action_allowed(service=svc, operation=name) - ) - for svc in services - } - return json.dumps(listing) - - def _describe_operation(self, service: str, operation: str) -> str: - canonical, model = self._resolve_allowed_operation(service, operation) - op = model.operation_model(canonical) - return json.dumps( - { - "service": service, - "operation": canonical, - "documentation": _strip_html(op.documentation)[:600], - "input": _describe_shape(op.input_shape), - "output": _describe_shape(op.output_shape), - } - ) - - def _call_aws(self, service: str, operation: str, parameters: dict[str, Any]) -> str: - canonical, _ = self._resolve_allowed_operation(service, operation) - client = self._get_client(service) - method_name = xform_name(canonical) - if client.can_paginate(method_name): - pages = client.get_paginator(method_name).paginate( - **parameters, PaginationConfig={"MaxItems": self._max_items} - ) - response = pages.build_full_result() - else: - response = getattr(client, method_name)(**parameters) - return self._serialize_response(response) - - # ------------------------------------------------------------------ - # Lazy client resolution and serialization - # ------------------------------------------------------------------ - - def _get_client(self, service: str) -> BaseClient: - if service not in self._clients: - hook = AwsBaseHook( - aws_conn_id=self._aws_conn_id, client_type=service, region_name=self._region_name - ) - self._clients[service] = hook.get_conn() - return self._clients[service] - - def _serialize_response(self, response: Any) -> str: - if isinstance(response, dict): - response.pop("ResponseMetadata", None) - payload = json.dumps(_to_jsonable(response, max_bytes=self._max_output_bytes), default=str) - if len(payload) > self._max_output_bytes: - return json.dumps( - { - "truncated": True, - "max_output_bytes": self._max_output_bytes, - "data": payload[: self._max_output_bytes], - } - ) - return payload - - -# --------------------------------------------------------------------------- -# Private botocore helpers -# --------------------------------------------------------------------------- - - -@cache -def _get_available_services() -> frozenset[str]: - return frozenset(botocore.session.get_session().get_available_services()) - - -@cache -def _load_service_model(service: str) -> ServiceModel: - return botocore.session.get_session().get_service_model(service) - - -def _resolve_operation_name(model: ServiceModel, operation: str) -> str | None: - """Resolve a case/underscore-insensitive operation name to its canonical API name.""" - wanted = _normalize_action(operation) - for name in model.operation_names: - if _normalize_action(name) == wanted: - return name - return None - - -def _describe_shape(shape: Shape | None, depth: int = _SHAPE_DEPTH) -> Any: - """Render a botocore shape as a compact JSON-friendly summary.""" - if shape is None: - return None - type_name = shape.type_name - if type_name == "structure": - if depth <= 0: - return type_name - required = set(shape.required_members) - return { - name: ( - {"type": _describe_shape(member, depth - 1), "required": True} - if name in required - else {"type": _describe_shape(member, depth - 1)} - ) - for name, member in shape.members.items() - } - if type_name == "list": - return [_describe_shape(shape.member, depth - 1)] if depth > 0 else type_name - if type_name == "map": - return {"": _describe_shape(shape.value, depth - 1)} if depth > 0 else type_name - enum = getattr(shape, "enum", None) - if enum: - return f"{type_name} (one of: {', '.join(enum)})" - return type_name - - -_HTML_TAG_RE = re.compile(r"<[^>]+>") - - -def _strip_html(text: str | None) -> str: - return " ".join(_HTML_TAG_RE.sub(" ", text or "").split()) - - -def _to_jsonable(value: Any, *, max_bytes: int) -> Any: - """Make a boto3 response JSON-serializable; datetimes are handled by ``default=str``.""" - if isinstance(value, StreamingBody): - # Read one byte past the cap so truncation is detectable. - raw = value.read(max_bytes + 1) - text = raw[:max_bytes].decode("utf-8", errors="replace") - if len(raw) > max_bytes: - return {"truncated": True, "data": text} - return text - if isinstance(value, bytes): - return value.decode("utf-8", errors="replace") - if isinstance(value, dict): - return {k: _to_jsonable(v, max_bytes=max_bytes) for k, v in value.items()} - if isinstance(value, list): - return [_to_jsonable(v, max_bytes=max_bytes) for v in value] - return value diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_aws.py b/providers/common/ai/tests/unit/common/ai/toolsets/test_aws.py deleted file mode 100644 index 08d29f12b90b1..0000000000000 --- a/providers/common/ai/tests/unit/common/ai/toolsets/test_aws.py +++ /dev/null @@ -1,359 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -from __future__ import annotations - -import asyncio -import json -from datetime import datetime, timezone -from unittest.mock import MagicMock, patch - -import pytest -from botocore.exceptions import ClientError -from pydantic_ai.exceptions import ModelRetry - -from airflow.providers.common.ai.toolsets.aws import ( - AWSToolset, - _get_available_services, - _load_service_model, - _resolve_operation_name, -) -from airflow.providers.common.ai.utils.tool_definition import _SUPPORTS_RETURN_SCHEMA - - -def _make_toolset(**kwargs): - kwargs.setdefault("allowed_actions", ["s3:List*", "s3:GetBucketLocation"]) - return AWSToolset("aws_test", **kwargs) - - -def _make_mock_client(): - client = MagicMock() - client.can_paginate.return_value = False - return client - - -def _call(ts, name, args): - return asyncio.run(ts.call_tool(name, args, ctx=MagicMock(), tool=MagicMock())) - - -class TestAWSToolsetInit: - def test_id_includes_conn_id(self): - ts = _make_toolset() - assert ts.id == "aws-aws_test" - - def test_rejects_empty_allowed_actions(self): - with pytest.raises(ValueError, match="non-empty"): - AWSToolset("aws_test", allowed_actions=[]) - - @pytest.mark.parametrize( - ("action", "match"), - [ - ("s3", "expected"), - ("s3:", "expected"), - (":ListBuckets", "expected"), - ("*:GetObject", "wildcards are not allowed in the service part"), - ("nosuchservice:DoThing", "Unknown AWS service"), - ("s3:NoSuchOperation", "Unknown operation"), - ], - ) - def test_rejects_invalid_action(self, action, match): - with pytest.raises(ValueError, match=match): - AWSToolset("aws_test", allowed_actions=[action]) - - def test_accepts_snake_case_operation(self): - AWSToolset("aws_test", allowed_actions=["s3:list_buckets"]) - - def test_rejects_wildcard_action_that_matches_no_operations(self): - with pytest.raises(ValueError, match="does not match any"): - AWSToolset("aws_test", allowed_actions=["s3:ListObjectz*"]) - - def test_rejects_wildcard_action_that_only_matches_sensitive_operations(self): - with pytest.raises(ValueError, match="does not match any"): - AWSToolset("aws_test", allowed_actions=["kms:Decrypt*"]) - - @patch("airflow.providers.common.ai.toolsets.aws._get_available_services", autospec=True) - def test_rejects_service_missing_from_available_services(self, mock_get_available_services): - mock_get_available_services.return_value = frozenset({"ec2"}) - - with pytest.raises(ValueError, match="Unknown AWS service 's3'"): - AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"]) - - -class TestAWSToolsetGetTools: - def test_returns_three_tools(self): - ts = _make_toolset() - tools = asyncio.run(ts.get_tools(ctx=MagicMock())) - assert set(tools.keys()) == {"list_aws_operations", "describe_aws_operation", "call_aws"} - - def test_tool_definitions_have_descriptions(self): - ts = _make_toolset() - tools = asyncio.run(ts.get_tools(ctx=MagicMock())) - for tool in tools.values(): - assert tool.tool_def.description - - @pytest.mark.skipif( - not _SUPPORTS_RETURN_SCHEMA, reason="pydantic-ai too old for ToolDefinition.return_schema" - ) - def test_tools_declare_string_return_schema(self): - ts = _make_toolset() - tools = asyncio.run(ts.get_tools(ctx=MagicMock())) - for tool in tools.values(): - assert tool.tool_def.return_schema == {"type": "string"} - - -class TestAWSToolsetActionMatching: - @pytest.mark.parametrize( - ("actions", "service", "operation", "allowed"), - [ - (["s3:ListBuckets"], "s3", "ListBuckets", True), - (["s3:list_buckets"], "s3", "ListBuckets", True), - (["s3:List*"], "s3", "ListObjectsV2", True), - (["s3:List*"], "s3", "GetObject", False), - (["s3:*"], "ec2", "DescribeInstances", False), - (["secretsmanager:*"], "secretsmanager", "GetSecretValue", False), - (["secretsmanager:GetSecretValue"], "secretsmanager", "GetSecretValue", True), - (["kms:*"], "kms", "Decrypt", False), - (["kms:Decrypt"], "kms", "Decrypt", True), - ], - ) - def test_allow_list_matching(self, actions, service, operation, allowed): - ts = AWSToolset("aws_test", allowed_actions=actions) - assert ts._is_action_allowed(service=service, operation=operation) is allowed - - -class TestAWSToolsetPrivateHelpers: - def test_is_action_allowed_requires_keyword_arguments(self): - ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"]) - - with pytest.raises(TypeError, match="positional"): - ts._is_action_allowed("s3", "ListBuckets") # type: ignore[misc] - - def test_pattern_match_allows_normal_operations_but_not_sensitive_operations(self): - ts = AWSToolset("aws_test", allowed_actions=["s3:List*", "kms:List*"]) - - assert ts._is_pattern_match_allowed("s3:list*", "s3", "ListBuckets") is True - assert ts._is_pattern_match_allowed("kms:decrypt*", "kms", "Decrypt") is False - - def test_resolve_allowed_service_returns_configured_service(self): - ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"]) - - assert ts._resolve_allowed_service("s3") == "s3" - - def test_resolve_allowed_service_rejects_service_outside_allow_list(self): - ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"]) - - with pytest.raises(ValueError, match="Service 'ec2' is not in this toolset's allowed actions"): - ts._resolve_allowed_service("ec2") - - def test_resolve_allowed_operation_returns_canonical_name_and_model(self): - ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"]) - - operation, model = ts._resolve_allowed_operation("s3", "list_buckets") - - assert operation == "ListBuckets" - assert "ListBuckets" in model.operation_names - - def test_resolve_allowed_operation_rejects_unknown_operation(self): - ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"]) - - with pytest.raises(ValueError, match="Unknown operation 'NoSuchOperation' for service 's3'"): - ts._resolve_allowed_operation("s3", "NoSuchOperation") - - def test_resolve_allowed_operation_rejects_disallowed_operation(self): - ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"]) - - with pytest.raises( - ValueError, match="Operation s3:GetObject is not in this toolset's allowed actions" - ): - ts._resolve_allowed_operation("s3", "GetObject") - - def test_resolve_allowed_operation_rejects_service_outside_allow_list(self): - ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"]) - - with pytest.raises(ValueError, match="Service 'ec2' is not in this toolset's allowed actions"): - ts._resolve_allowed_operation("ec2", "DescribeInstances") - - def test_get_available_services_reads_botocore_catalog(self): - assert "s3" in _get_available_services() - - def test_load_service_model_and_resolve_operation_name(self): - model = _load_service_model("s3") - - assert _resolve_operation_name(model, "list_buckets") == "ListBuckets" - assert _resolve_operation_name(model, "NoSuchOperation") is None - - -class TestAWSToolsetListOperations: - def test_lists_only_allowed_operations(self): - ts = AWSToolset("aws_test", allowed_actions=["s3:List*"]) - listing = json.loads(_call(ts, "list_aws_operations", {})) - assert "ListBuckets" in listing["s3"] - assert "GetObject" not in listing["s3"] - - def test_sensitive_operations_hidden_behind_wildcard(self): - ts = AWSToolset("aws_test", allowed_actions=["secretsmanager:*"]) - listing = json.loads(_call(ts, "list_aws_operations", {})) - assert "ListSecrets" in listing["secretsmanager"] - assert "GetSecretValue" not in listing["secretsmanager"] - - def test_verbatim_sensitive_operation_is_listed(self): - ts = AWSToolset("aws_test", allowed_actions=["secretsmanager:*", "secretsmanager:GetSecretValue"]) - listing = json.loads(_call(ts, "list_aws_operations", {})) - assert "GetSecretValue" in listing["secretsmanager"] - - def test_service_outside_allow_list_raises_model_retry(self): - ts = _make_toolset() - with pytest.raises(ModelRetry, match="not in this toolset's allowed actions"): - _call(ts, "list_aws_operations", {"service": "ec2"}) - - -class TestAWSToolsetDescribeOperation: - def test_returns_input_shape_with_required_members(self): - ts = AWSToolset("aws_test", allowed_actions=["s3:ListObjectsV2"]) - described = json.loads( - _call(ts, "describe_aws_operation", {"service": "s3", "operation": "ListObjectsV2"}) - ) - assert described["operation"] == "ListObjectsV2" - assert described["input"]["Bucket"]["required"] is True - assert "Contents" in described["output"] - - def test_resolves_snake_case_operation_name(self): - ts = AWSToolset("aws_test", allowed_actions=["s3:ListObjectsV2"]) - described = json.loads( - _call(ts, "describe_aws_operation", {"service": "s3", "operation": "list_objects_v2"}) - ) - assert described["operation"] == "ListObjectsV2" - - @pytest.mark.parametrize( - ("service", "operation", "match"), - [ - ("s3", "GetObject", "not in this toolset's allowed actions"), - ("ec2", "DescribeInstances", "Service 'ec2' is not in this toolset's allowed actions"), - ("s3", "NoSuchOperation", "Unknown operation 'NoSuchOperation' for service 's3'"), - ], - ids=["disallowed_operation", "service_outside_allow_list", "unknown_operation"], - ) - def test_model_supplied_invalid_operation_raises_model_retry(self, service, operation, match): - ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"]) - - with pytest.raises(ModelRetry, match=match): - _call(ts, "describe_aws_operation", {"service": service, "operation": operation}) - - -class TestAWSToolsetCallAws: - def test_executes_operation_and_returns_json(self): - ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"]) - client = _make_mock_client() - client.list_buckets.return_value = { - "Buckets": [ - { - "Name": "data-lake-raw", - "CreationDate": datetime(2026, 1, 1, tzinfo=timezone.utc), - } - ], - "ResponseMetadata": {"RequestId": "abc123"}, - } - ts._clients["s3"] = client - - result = json.loads(_call(ts, "call_aws", {"service": "s3", "operation": "ListBuckets"})) - client.list_buckets.assert_called_once_with() - assert result["Buckets"][0]["Name"] == "data-lake-raw" - # datetime serialized via default=str, ResponseMetadata stripped. - assert "2026-01-01" in result["Buckets"][0]["CreationDate"] - assert "ResponseMetadata" not in result - - def test_snake_case_operation_resolves(self): - ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"]) - client = _make_mock_client() - client.list_buckets.return_value = {"Buckets": []} - ts._clients["s3"] = client - - result = json.loads(_call(ts, "call_aws", {"service": "s3", "operation": "list_buckets"})) - assert result == {"Buckets": []} - - def test_uses_paginator_when_available(self): - ts = AWSToolset("aws_test", allowed_actions=["s3:ListObjectsV2"], max_items=250) - client = _make_mock_client() - client.can_paginate.return_value = True - paginator = client.get_paginator.return_value - paginator.paginate.return_value.build_full_result.return_value = {"Contents": [{"Key": "a.parquet"}]} - ts._clients["s3"] = client - - result = json.loads( - _call( - ts, - "call_aws", - {"service": "s3", "operation": "ListObjectsV2", "parameters": {"Bucket": "b"}}, - ) - ) - paginator.paginate.assert_called_once_with(Bucket="b", PaginationConfig={"MaxItems": 250}) - assert result["Contents"] == [{"Key": "a.parquet"}] - - @pytest.mark.parametrize( - ("service", "operation", "match"), - [ - ("s3", "DeleteBucket", "not in this toolset's allowed actions"), - ("ec2", "DescribeInstances", "Service 'ec2' is not in this toolset's allowed actions"), - ("s3", "NoSuchOperation", "Unknown operation 'NoSuchOperation' for service 's3'"), - ], - ids=["disallowed_operation", "service_outside_allow_list", "unknown_operation"], - ) - def test_model_supplied_invalid_operation_raises_model_retry_without_client( - self, service, operation, match - ): - ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"]) - - with pytest.raises(ModelRetry, match=match): - _call(ts, "call_aws", {"service": service, "operation": operation}) - assert ts._clients == {} - - def test_client_error_raises_model_retry(self): - ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"]) - client = _make_mock_client() - client.list_buckets.side_effect = ClientError( - {"Error": {"Code": "AccessDenied", "Message": "not authorized"}}, "ListBuckets" - ) - ts._clients["s3"] = client - - with pytest.raises(ModelRetry, match="AccessDenied"): - _call(ts, "call_aws", {"service": "s3", "operation": "ListBuckets"}) - - def test_large_response_is_truncated(self): - ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"], max_output_bytes=200) - client = _make_mock_client() - client.list_buckets.return_value = {"Buckets": [{"Name": "x" * 50}] * 50} - ts._clients["s3"] = client - - result = json.loads(_call(ts, "call_aws", {"service": "s3", "operation": "ListBuckets"})) - assert result["truncated"] is True - assert len(result["data"]) == 200 - - @patch("airflow.providers.common.ai.toolsets.aws.AwsBaseHook", autospec=True) - def test_client_resolved_via_amazon_hook(self, mock_hook_cls): - client = _make_mock_client() - client.list_buckets.return_value = {"Buckets": []} - mock_hook_cls.return_value.get_conn.return_value = client - - ts = AWSToolset("aws_prod", allowed_actions=["s3:ListBuckets"], region_name="eu-west-1") - _call(ts, "call_aws", {"service": "s3", "operation": "ListBuckets"}) - mock_hook_cls.assert_called_once_with( - aws_conn_id="aws_prod", client_type="s3", region_name="eu-west-1" - ) - - def test_unknown_tool_raises_value_error(self): - ts = _make_toolset() - with pytest.raises(ValueError, match="Unknown tool"): - _call(ts, "use_aws", {}) diff --git a/uv.lock b/uv.lock index 8bc48d4c7af97..41261e1ec47d0 100644 --- a/uv.lock +++ b/uv.lock @@ -4396,18 +4396,12 @@ dependencies = [ ] [package.optional-dependencies] -amazon = [ - { name = "apache-airflow-providers-amazon" }, -] anthropic = [ { name = "pydantic-ai-slim", extra = ["anthropic"] }, ] avro = [ { name = "fastavro" }, ] -aws = [ - { name = "apache-airflow-providers-amazon" }, -] bedrock = [ { name = "pydantic-ai-slim", extra = ["bedrock"] }, ] @@ -4463,7 +4457,6 @@ sql = [ dev = [ { name = "apache-airflow" }, { name = "apache-airflow-devel-common" }, - { name = "apache-airflow-providers-amazon" }, { name = "apache-airflow-providers-common-compat" }, { name = "apache-airflow-providers-common-sql", extra = ["datafusion"] }, { name = "apache-airflow-providers-git" }, @@ -4484,8 +4477,6 @@ docs = [ [package.metadata] requires-dist = [ { name = "apache-airflow", editable = "." }, - { name = "apache-airflow-providers-amazon", marker = "extra == 'amazon'", editable = "providers/amazon" }, - { name = "apache-airflow-providers-amazon", marker = "extra == 'aws'", editable = "providers/amazon" }, { name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" }, { name = "apache-airflow-providers-common-sql", marker = "extra == 'common-sql'", editable = "providers/common/sql" }, { name = "apache-airflow-providers-common-sql", marker = "extra == 'sql'", editable = "providers/common/sql" }, @@ -4514,13 +4505,12 @@ requires-dist = [ { name = "python-docx", marker = "extra == 'docx'", specifier = ">=1.0.0" }, { name = "sqlglot", marker = "extra == 'sql'", specifier = ">=30.0.0" }, ] -provides-extras = ["anthropic", "bedrock", "google", "openai", "mcp", "code-mode", "shields", "skills", "avro", "parquet", "sql", "aws", "common-sql", "langchain", "llamaindex", "pdf", "docx", "git", "amazon"] +provides-extras = ["anthropic", "bedrock", "google", "openai", "mcp", "code-mode", "shields", "skills", "avro", "parquet", "sql", "common-sql", "langchain", "llamaindex", "pdf", "docx", "git"] [package.metadata.requires-dev] dev = [ { name = "apache-airflow", editable = "." }, { name = "apache-airflow-devel-common", editable = "devel-common" }, - { name = "apache-airflow-providers-amazon", editable = "providers/amazon" }, { name = "apache-airflow-providers-common-compat", editable = "providers/common/compat" }, { name = "apache-airflow-providers-common-sql", editable = "providers/common/sql" }, { name = "apache-airflow-providers-common-sql", extras = ["datafusion"], editable = "providers/common/sql" }, From 012d5c6ec14f82bf4ff4e5dd5d85e0b20b85212c Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Thu, 30 Jul 2026 01:38:38 +0530 Subject: [PATCH 2/3] Remove AWS toolset module from common.ai provider.yaml The AWS toolset module was deleted by this revert but stayed registered under `toolsets/python-modules` in provider.yaml, so the provider.yaml validation static check failed importing it. --- providers/common/ai/provider.yaml | 1 - 1 file changed, 1 deletion(-) diff --git a/providers/common/ai/provider.yaml b/providers/common/ai/provider.yaml index 02d69094da2a0..181696f12b461 100644 --- a/providers/common/ai/provider.yaml +++ b/providers/common/ai/provider.yaml @@ -455,7 +455,6 @@ operators: toolsets: - integration-name: Common AI python-modules: - - airflow.providers.common.ai.toolsets.aws - airflow.providers.common.ai.toolsets.hook - airflow.providers.common.ai.toolsets.sql - airflow.providers.common.ai.toolsets.datafusion From b197d55d0eb23a519d026e7925ccdf2b122771d3 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Thu, 30 Jul 2026 01:39:34 +0530 Subject: [PATCH 3/3] Mark common.ai provider as ready for release Reverses #70481. That marked the provider not-ready so the release manager would skip it while the AWS toolset work was still in flight; this PR removes that toolset, so the provider is releasable again. Flipping the state back re-adds common.ai to the provider-id choice lists in the release-management help screens, hence the regenerated breeze command output images. --- ...se-management_classify-provider-changes.svg | 16 ++++++++-------- ...se-management_classify-provider-changes.txt | 2 +- ...gement_generate-issue-content-providers.svg | 16 ++++++++-------- ...gement_generate-issue-content-providers.txt | 2 +- ...-management_generate-providers-metadata.svg | 18 +++++++++--------- ...-management_generate-providers-metadata.txt | 2 +- ...nagement_prepare-provider-distributions.svg | 16 ++++++++-------- ...nagement_prepare-provider-distributions.txt | 2 +- ...nagement_prepare-provider-documentation.svg | 16 ++++++++-------- ...nagement_prepare-provider-documentation.txt | 2 +- providers/common/ai/provider.yaml | 2 +- 11 files changed, 47 insertions(+), 47 deletions(-) diff --git a/dev/breeze/doc/images/output_release-management_classify-provider-changes.svg b/dev/breeze/doc/images/output_release-management_classify-provider-changes.svg index 53946097209cd..cef6e4d3b1d14 100644 --- a/dev/breeze/doc/images/output_release-management_classify-provider-changes.svg +++ b/dev/breeze/doc/images/output_release-management_classify-provider-changes.svg @@ -165,14 +165,14 @@ [OPTIONS] [airbyte | akeyless | alibaba | amazon | anthropic | apache.cassandra | apache.drill | apache.druid |        apache.flink | apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy | apache.pig | apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | atlassian.jira | celery | -clickhousedb | cloudant | cncf.kubernetes | cohere | common.compat | common.io | common.messaging | common.sql |       -databricks | datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab | facebook | ftp -| git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins | keycloak |        -microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql | neo4j | odbc | openai |         -openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | papermill | pgvector | pinecone | postgres |     -presto | qdrant | redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp | snowflake |    -sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex | ydb |          -zendesk]...                                                                                                            +clickhousedb | cloudant | cncf.kubernetes | cohere | common.ai | common.compat | common.io | common.messaging |        +common.sql | databricks | datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab |   +facebook | ftp | git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins |    +keycloak | microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql | neo4j | odbc |       +openai | openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | papermill | pgvector | pinecone |       +postgres | presto | qdrant | redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp |     +snowflake | sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex |    +ydb | zendesk]...                                                                                                      Classify each provider's unreleased changes with hard-coded, high-confidence rules, flagging ambiguous commits as  'needs_llm' for an agent/skill to assess. Outputs JSON - a deterministic alternative to the random '--non-interactive' diff --git a/dev/breeze/doc/images/output_release-management_classify-provider-changes.txt b/dev/breeze/doc/images/output_release-management_classify-provider-changes.txt index 738d0ea6511eb..6e02c700afdf5 100644 --- a/dev/breeze/doc/images/output_release-management_classify-provider-changes.txt +++ b/dev/breeze/doc/images/output_release-management_classify-provider-changes.txt @@ -1 +1 @@ -3d2a0b7552829f49477890b293207a4c +4cac13b21eee8b732a46c5a15aec7a4b diff --git a/dev/breeze/doc/images/output_release-management_generate-issue-content-providers.svg b/dev/breeze/doc/images/output_release-management_generate-issue-content-providers.svg index 843f40c13e3ef..2f9692be54fa0 100644 --- a/dev/breeze/doc/images/output_release-management_generate-issue-content-providers.svg +++ b/dev/breeze/doc/images/output_release-management_generate-issue-content-providers.svg @@ -154,14 +154,14 @@ [OPTIONS] [airbyte | akeyless | alibaba | amazon | anthropic | apache.cassandra | apache.drill | apache.druid |        apache.flink | apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy | apache.pig | apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | atlassian.jira | celery | -clickhousedb | cloudant | cncf.kubernetes | cohere | common.compat | common.io | common.messaging | common.sql |       -databricks | datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab | facebook | ftp -| git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins | keycloak |        -microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql | neo4j | odbc | openai |         -openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | papermill | pgvector | pinecone | postgres |     -presto | qdrant | redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp | snowflake |    -sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex | ydb |          -zendesk]...                                                                                                            +clickhousedb | cloudant | cncf.kubernetes | cohere | common.ai | common.compat | common.io | common.messaging |        +common.sql | databricks | datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab |   +facebook | ftp | git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins |    +keycloak | microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql | neo4j | odbc |       +openai | openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | papermill | pgvector | pinecone |       +postgres | presto | qdrant | redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp |     +snowflake | sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex |    +ydb | zendesk]...                                                                                                      Generates content for issue to test the release. diff --git a/dev/breeze/doc/images/output_release-management_generate-issue-content-providers.txt b/dev/breeze/doc/images/output_release-management_generate-issue-content-providers.txt index e8dc57000663a..6afb289b5854f 100644 --- a/dev/breeze/doc/images/output_release-management_generate-issue-content-providers.txt +++ b/dev/breeze/doc/images/output_release-management_generate-issue-content-providers.txt @@ -1 +1 @@ -abd6c12fe888796a153522495b802278 +bd80bf2dd63a27111b8c5745df5dfb86 diff --git a/dev/breeze/doc/images/output_release-management_generate-providers-metadata.svg b/dev/breeze/doc/images/output_release-management_generate-providers-metadata.svg index 8df40d73a499b..64efb0af6b65b 100644 --- a/dev/breeze/doc/images/output_release-management_generate-providers-metadata.svg +++ b/dev/breeze/doc/images/output_release-management_generate-providers-metadata.svg @@ -173,15 +173,15 @@ │apache.flink | apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | │ │apache.kylin | apache.livy | apache.pig | apache.pinot | apache.spark | apache.tinkerpop | │ │apprise | arangodb | asana | atlassian.jira | celery | clickhousedb | cloudant | cncf.kubernetes│ -│| cohere | common.compat | common.io | common.messaging | common.sql | databricks | datadog | │ -│dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab | facebook | ftp │ -│| git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | │ -│jenkins | keycloak | microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | │ -│mongo | mysql | neo4j | odbc | openai | openfaas | openlineage | opensearch | opsgenie | oracle │ -│| pagerduty | papermill | pgvector | pinecone | postgres | presto | qdrant | redis | salesforce │ -│| samba | segment | sendgrid | sftp | singularity | slack | smtp | snowflake | sqlite | ssh | │ -│standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex | ydb | │ -│zendesk)│ +│| cohere | common.ai | common.compat | common.io | common.messaging | common.sql | databricks | │ +│datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab | │ +│facebook | ftp | git | github | google | grpc | hashicorp | http | imap | influxdb | informatica│ +│| jdbc | jenkins | keycloak | microsoft.azure | microsoft.mssql | microsoft.psrp | │ +│microsoft.winrm | mongo | mysql | neo4j | odbc | openai | openfaas | openlineage | opensearch | │ +│opsgenie | oracle | pagerduty | papermill | pgvector | pinecone | postgres | presto | qdrant | │ +│redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp | snowflake │ +│| sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate |│ +│yandex | ydb | zendesk)│ │--provider-versionProvider version to generate metadata for. Only used when --provider-id is specified. Limits    │ │running metadata generation to only this version of the provider. (TEXT)│ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ diff --git a/dev/breeze/doc/images/output_release-management_generate-providers-metadata.txt b/dev/breeze/doc/images/output_release-management_generate-providers-metadata.txt index 153ac6f7b0699..7c14e04e63151 100644 --- a/dev/breeze/doc/images/output_release-management_generate-providers-metadata.txt +++ b/dev/breeze/doc/images/output_release-management_generate-providers-metadata.txt @@ -1 +1 @@ -010c3a7215ff699933b21bf4e5787fcb +5865810f5a4b12c9199267b7dd6459b3 diff --git a/dev/breeze/doc/images/output_release-management_prepare-provider-distributions.svg b/dev/breeze/doc/images/output_release-management_prepare-provider-distributions.svg index f6acf9ed8295e..27a1d92de9a4b 100644 --- a/dev/breeze/doc/images/output_release-management_prepare-provider-distributions.svg +++ b/dev/breeze/doc/images/output_release-management_prepare-provider-distributions.svg @@ -192,14 +192,14 @@ [OPTIONS] [airbyte | akeyless | alibaba | amazon | anthropic | apache.cassandra | apache.drill | apache.druid |        apache.flink | apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy | apache.pig | apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | atlassian.jira | celery | -clickhousedb | cloudant | cncf.kubernetes | cohere | common.compat | common.io | common.messaging | common.sql |       -databricks | datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab | facebook | ftp -| git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins | keycloak |        -microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql | neo4j | odbc | openai |         -openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | papermill | pgvector | pinecone | postgres |     -presto | qdrant | redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp | snowflake |    -sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex | ydb |          -zendesk]...                                                                                                            +clickhousedb | cloudant | cncf.kubernetes | cohere | common.ai | common.compat | common.io | common.messaging |        +common.sql | databricks | datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab |   +facebook | ftp | git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins |    +keycloak | microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql | neo4j | odbc |       +openai | openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | papermill | pgvector | pinecone |       +postgres | presto | qdrant | redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp |     +snowflake | sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex |    +ydb | zendesk]...                                                                                                      Prepare sdist/whl distributions of Airflow Providers. Each provider directory is wiped with `git clean -fdx`  (preserving .venv, .idea, .vscode) before build to keep in-tree generated files out of the artifact. See dev/breeze  diff --git a/dev/breeze/doc/images/output_release-management_prepare-provider-distributions.txt b/dev/breeze/doc/images/output_release-management_prepare-provider-distributions.txt index f71d10374df6c..f054c03806e05 100644 --- a/dev/breeze/doc/images/output_release-management_prepare-provider-distributions.txt +++ b/dev/breeze/doc/images/output_release-management_prepare-provider-distributions.txt @@ -1 +1 @@ -e8af9b706de3926a3297a6c1114d7e5c +a27c1726f5902e5fdb501ecdee226476 diff --git a/dev/breeze/doc/images/output_release-management_prepare-provider-documentation.svg b/dev/breeze/doc/images/output_release-management_prepare-provider-documentation.svg index 32d59719218eb..b9c71f9c81aec 100644 --- a/dev/breeze/doc/images/output_release-management_prepare-provider-documentation.svg +++ b/dev/breeze/doc/images/output_release-management_prepare-provider-documentation.svg @@ -213,14 +213,14 @@ [OPTIONS] [airbyte | akeyless | alibaba | amazon | anthropic | apache.cassandra | apache.drill | apache.druid |        apache.flink | apache.hdfs | apache.hive | apache.iceberg | apache.impala | apache.kafka | apache.kylin | apache.livy | apache.pig | apache.pinot | apache.spark | apache.tinkerpop | apprise | arangodb | asana | atlassian.jira | celery | -clickhousedb | cloudant | cncf.kubernetes | cohere | common.compat | common.io | common.messaging | common.sql |       -databricks | datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab | facebook | ftp -| git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins | keycloak |        -microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql | neo4j | odbc | openai |         -openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | papermill | pgvector | pinecone | postgres |     -presto | qdrant | redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp | snowflake |    -sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex | ydb |          -zendesk]...                                                                                                            +clickhousedb | cloudant | cncf.kubernetes | cohere | common.ai | common.compat | common.io | common.messaging |        +common.sql | databricks | datadog | dbt.cloud | dingding | discord | docker | edge3 | elasticsearch | exasol | fab |   +facebook | ftp | git | github | google | grpc | hashicorp | http | imap | influxdb | informatica | jdbc | jenkins |    +keycloak | microsoft.azure | microsoft.mssql | microsoft.psrp | microsoft.winrm | mongo | mysql | neo4j | odbc |       +openai | openfaas | openlineage | opensearch | opsgenie | oracle | pagerduty | papermill | pgvector | pinecone |       +postgres | presto | qdrant | redis | salesforce | samba | segment | sendgrid | sftp | singularity | slack | smtp |     +snowflake | sqlite | ssh | standard | tableau | telegram | teradata | trino | vertica | vespa | weaviate | yandex |    +ydb | zendesk]...                                                                                                      Prepare CHANGELOG, README and COMMITS information for providers. diff --git a/dev/breeze/doc/images/output_release-management_prepare-provider-documentation.txt b/dev/breeze/doc/images/output_release-management_prepare-provider-documentation.txt index e4e1fc0ee6501..a54080b6bfbb5 100644 --- a/dev/breeze/doc/images/output_release-management_prepare-provider-documentation.txt +++ b/dev/breeze/doc/images/output_release-management_prepare-provider-documentation.txt @@ -1 +1 @@ -33f340084282725cf473d84c59e2d63f +c33a2f6d00a3a8dbec8b56c1c2d88d54 diff --git a/providers/common/ai/provider.yaml b/providers/common/ai/provider.yaml index 181696f12b461..7431bcde2c6c3 100644 --- a/providers/common/ai/provider.yaml +++ b/providers/common/ai/provider.yaml @@ -21,7 +21,7 @@ name: Common AI description: | AI/LLM hooks and operators for Airflow pipelines using `pydantic-ai `__. -state: not-ready +state: ready lifecycle: incubation source-date-epoch: 1784746595 build-system: hatchling