This is an automated email from the ASF dual-hosted git repository.
kaxil 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 c10185f281d Fix PydanticAI vendor connections not resolving from
secrets backends (#72853)
c10185f281d is described below
commit c10185f281d2b91b6a42dc1b2abdd4cbe9e4f217
Author: Kaxil Naik <[email protected]>
AuthorDate: Thu Sep 10 22:25:46 2026 +0100
Fix PydanticAI vendor connections not resolving from secrets backends
(#72853)
* Fix PydanticAI vendor connections not resolving from secrets backends
The three vendor connection types declared their `conn_type` with hyphens
(`pydanticai-azure`, `pydanticai-bedrock`, `pydanticai-vertex`), but a hook
is
registered under the literal `connection-type` from `provider.yaml` while
`Connection.from_uri` and `Connection.from_json` rewrite `-` to `_` before
the
lookup. The registered key was therefore unreachable from every secrets
backend,
including the `AIRFLOW_CONN_*` environment variables enabled by default,
and any
task using one of these connections failed with:
AirflowException: Unknown hook type "pydanticai_azure"
Only a connection read straight out of the metadata DB kept its hyphen and
resolved, which is why this went unnoticed.
Renames the three types to underscores, matching every other multi-word
hook in
the provider tree (`azure_data_factory`, `google_cloud_platform`,
`spark_connect`).
`default_conn_name` values were already underscored and are unchanged.
* Reject hyphenated connection-type in the provider schema and document the
rename
Adds `"pattern": "^[a-z0-9_]+$"` to the connection-type property in both
provider
schemas, so a hyphenated type fails schema validation at authoring time
instead of
at a user's first task run. All 141 connection-type values across the
provider tree
already satisfy it.
Also parametrizes the round-trip guard so it reports every offending type
rather
than aborting at the first, and notes the rename in the three connection
docs for
users whose connections are already stored with the old type.
* Keep the connection-type gate in the authoring schema only
The runtime schema (provider_info.schema.json) is validated against the
get_provider_info() output of every installed provider, third-party
included, and
that validate() call is unwrapped in provider discovery. A pattern there
turns one
non-conforming provider into a failure for all provider discovery in the
process,
and every released apache-airflow-providers-common-ai wheel declares the
hyphenated
types, so installing a released provider against new core would fail to
import
Airflow at all. PR #71104 declined the same tightening for the same reason.
The
authoring gate is fully delivered by provider.yaml.schema.json, which the
provider-yaml pre-commit validates for every in-tree provider.
Tightens that authoring pattern to ^[a-z][a-z0-9_]*$ so it also rejects a
leading
digit, which is not a valid URI scheme under RFC 3986 and which Python 3.10
and
3.14 parse differently.
Corrects the three connection-doc notes, which claimed every stored
connection
needed updating. URI and JSON values decode '-' back to '_' and resolve
unchanged,
so only types stored verbatim (metadata DB rows and object-form imports)
need a
change; those now carry the procedure. Also fixes the hyphenated
connection-type
example in the new-provider guide, and pins the resolution tests to the
exact hook
class with a per-parametrization conn_id.
* Anchor the connection-type pattern past a trailing newline
A schema pattern is matched with re.search, and Python's '$' matches
before a trailing newline as well as at the end of the string, so
^[a-z][a-z0-9_]*$ accepted a connection type carrying one. A YAML block
scalar produces exactly that, so 'connection-type: |' followed by an
indented name would have registered a hook under a name no connection can
reach.
A hyphen was already rejected either way, since '-' is outside the
character class and the newline allowance excuses only a newline, so this
does not reopen the case the pattern was added for. It closes the one
value that could still slip past it.
Adds the first test for this pattern, covering the spellings that cannot
be resolved: a hyphen, a block scalar's trailing newline, uppercase, a
leading digit, whitespace and an embedded newline.
---
airflow-core/src/airflow/provider.yaml.schema.json | 5 +-
.../tests/unit/always/test_provider_yaml_schema.py | 73 +++++++++++++++++
.../23_provider_hook_migration_to_yaml.rst | 2 +-
providers/MANAGING_PROVIDERS_LIFECYCLE.rst | 2 +-
.../ai/docs/connections/pydantic_ai_azure.rst | 29 ++++++-
.../ai/docs/connections/pydantic_ai_bedrock.rst | 33 ++++++--
.../ai/docs/connections/pydantic_ai_vertex.rst | 31 +++++++-
providers/common/ai/docs/index.rst | 2 +-
providers/common/ai/provider.yaml | 6 +-
.../providers/common/ai/get_provider_info.py | 6 +-
.../providers/common/ai/hooks/pydantic_ai.py | 6 +-
.../airflow/providers/common/ai/operators/llm.py | 2 +-
.../tests/unit/common/ai/hooks/test_pydantic_ai.py | 91 +++++++++++++++++++---
13 files changed, 249 insertions(+), 39 deletions(-)
diff --git a/airflow-core/src/airflow/provider.yaml.schema.json
b/airflow-core/src/airflow/provider.yaml.schema.json
index fd449394c18..4c4a199a882 100644
--- a/airflow-core/src/airflow/provider.yaml.schema.json
+++ b/airflow-core/src/airflow/provider.yaml.schema.json
@@ -428,8 +428,9 @@
"type": "object",
"properties": {
"connection-type": {
- "description": "Type of connection defined by the
provider",
- "type": "string"
+ "description": "Type of connection defined by the
provider. Lowercase letters, digits and '_' only, starting with a letter. The
hook is registered under this exact string, while Connection.get_uri() encodes
'_' as '-' and from_uri/from_json decode '-' back to '_'. A '-' here is
therefore indistinguishable from an encoded '_' and leaves the hook
unresolvable for any connection stored as a URI or JSON; uppercase is lost to
get_uri()'s lowercasing; and a leading digi [...]
+ "type": "string",
+ "pattern": "^[a-z][a-z0-9_]*$(?!\\n)"
},
"hook-class-name": {
"description": "Hook class name that implements the
connection type",
diff --git a/airflow-core/tests/unit/always/test_provider_yaml_schema.py
b/airflow-core/tests/unit/always/test_provider_yaml_schema.py
new file mode 100644
index 00000000000..70316be79e3
--- /dev/null
+++ b/airflow-core/tests/unit/always/test_provider_yaml_schema.py
@@ -0,0 +1,73 @@
+#
+# 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 json
+from pathlib import Path
+
+import jsonschema
+import pytest
+import yaml
+
+import airflow
+
+SCHEMA_PATH = Path(airflow.__file__).parent / "provider.yaml.schema.json"
+
+# A connection type is registered under this exact string, so the authoring
schema constrains
+# it to the one spelling that can be reached again: get_uri() lowercases and
encodes '_' as
+# '-', and reading a connection back decodes '-' to '_'.
+CONNECTION_TYPE_SCHEMA =
json.loads(SCHEMA_PATH.read_text())["properties"]["connection-types"]["items"][
+ "properties"
+]["connection-type"]
+
+# A YAML block scalar keeps a trailing newline, which is the value most easily
written by
+# accident. It is built here through the parser rather than hard-coded, so the
case documents
+# how such a value reaches the schema at all.
+BLOCK_SCALAR_VALUE = yaml.safe_load("connection-type: |\n
pydanticai_vertex\n")["connection-type"]
+
+
[email protected](
+ "connection_type",
+ ["pydanticai", "google_cloud_platform", "pagerduty_events", "a", "s3"],
+)
+def test_schema_accepts_a_registrable_connection_type(connection_type):
+ jsonschema.validate(connection_type, schema=CONNECTION_TYPE_SCHEMA)
+
+
[email protected](
+ ("connection_type", "reason"),
+ [
+ pytest.param("pydanticai-vertex", "'-' is the URI encoding of '_'",
id="hyphen"),
+ pytest.param(BLOCK_SCALAR_VALUE, "a trailing newline is not part of
the name", id="block-scalar"),
+ pytest.param("pydanticai-vertex\n", "hyphen and trailing newline",
id="hyphen-and-newline"),
+ pytest.param("PydanticAI", "get_uri() lowercases the scheme",
id="uppercase"),
+ pytest.param("1password", "a URI scheme cannot start with a digit",
id="leading-digit"),
+ pytest.param("pydantic ai", "a space cannot appear in a scheme",
id="space"),
+ pytest.param("pydantic\nai", "an embedded newline is not a scheme",
id="embedded-newline"),
+ pytest.param("", "a connection type is required to be non-empty",
id="empty"),
+ ],
+)
+def
test_schema_rejects_a_connection_type_that_cannot_be_resolved(connection_type,
reason):
+ """
+ The trailing-newline cases are why the pattern carries a lookahead rather
than ending at
+ '$'. A schema pattern is matched with ``re.search``, and Python's '$'
matches before a
+ trailing newline as well as at the end of the string, so
``^[a-z][a-z0-9_]*$`` alone
+ accepts a value written as a YAML block scalar.
+ """
+ with pytest.raises(jsonschema.ValidationError):
+ jsonschema.validate(connection_type, schema=CONNECTION_TYPE_SCHEMA)
diff --git a/contributing-docs/23_provider_hook_migration_to_yaml.rst
b/contributing-docs/23_provider_hook_migration_to_yaml.rst
index 54706f6c4dc..3b0069dbeb4 100644
--- a/contributing-docs/23_provider_hook_migration_to_yaml.rst
+++ b/contributing-docs/23_provider_hook_migration_to_yaml.rst
@@ -82,7 +82,7 @@ no effect on docs generation, logos, or tags.
connection-types:
- hook-class-name:
airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIAzureHook
hook-name: "Pydantic AI (Azure OpenAI)"
- connection-type: pydanticai-azure
+ connection-type: pydanticai_azure
external-services:
- Azure OpenAI
diff --git a/providers/MANAGING_PROVIDERS_LIFECYCLE.rst
b/providers/MANAGING_PROVIDERS_LIFECYCLE.rst
index b3d2076c004..08b000ea59b 100644
--- a/providers/MANAGING_PROVIDERS_LIFECYCLE.rst
+++ b/providers/MANAGING_PROVIDERS_LIFECYCLE.rst
@@ -270,7 +270,7 @@ your provider:
connection-types:
- hook-class-name:
airflow.providers.<PROVIDER>.hooks.<PROVIDER>.NewProviderHook
- - connection-type: provider-connection-type
+ connection-type: provider_connection_type
Building documentation locally
diff --git a/providers/common/ai/docs/connections/pydantic_ai_azure.rst
b/providers/common/ai/docs/connections/pydantic_ai_azure.rst
index 6660010d43a..4c8ee7daf65 100644
--- a/providers/common/ai/docs/connections/pydantic_ai_azure.rst
+++ b/providers/common/ai/docs/connections/pydantic_ai_azure.rst
@@ -15,18 +15,41 @@
specific language governing permissions and limitations
under the License.
-.. _howto/connection:pydanticai-azure:
+.. _howto/connection:pydanticai_azure:
Pydantic AI (Azure OpenAI) Connection
======================================
-The ``pydanticai-azure`` connection type configures access to
+The ``pydanticai_azure`` connection type configures access to
`Azure OpenAI
<https://azure.microsoft.com/en-us/products/ai-services/openai-service>`__
via the pydantic-ai framework. It backs ``PydanticAIAzureHook``, the dedicated
subclass of ``PydanticAIHook`` for Azure's non-standard auth (an endpoint URL
plus an API version, rather than the plain ``api_key`` + optional ``base_url``
that the generic :doc:`pydantic_ai` connection assumes).
+.. note::
+
+ This connection type was previously named ``pydanticai-azure``.
+
+ Connections stored as a URI or as JSON need no change: ``-`` is how ``_``
is
+ encoded in a URI scheme, so ``pydanticai-azure`` is decoded to
``pydanticai_azure``
+ on read and resolves as before. That covers ``AIRFLOW_CONN_*`` environment
+ variables and secrets backends such as HashiCorp Vault, AWS Secrets
Manager and
+ GCP Secret Manager.
+
+ A connection whose type is stored verbatim does need updating, because the
+ hyphen is preserved and no longer matches a registered hook. That means
rows in
+ the metadata database, including any created through the UI, and
connections
+ imported in object form from a local file:
+
+ .. code-block:: bash
+
+ airflow connections get <conn_id> -o json # confirm conn_type is
'pydanticai-azure'
+ airflow connections delete <conn_id>
+ airflow connections add <conn_id> --conn-type pydanticai_azure ...
+
+ In the UI, edit the connection and re-pick its type.
+
Default Connection IDs
----------------------
@@ -60,7 +83,7 @@ Examples
.. code-block:: json
{
- "conn_type": "pydanticai-azure",
+ "conn_type": "pydanticai_azure",
"password": "<azure-api-key>",
"host": "https://<resource>.openai.azure.com",
"extra": "{\"model\": \"azure:gpt-4o\", \"api_version\":
\"2024-07-01-preview\"}"
diff --git a/providers/common/ai/docs/connections/pydantic_ai_bedrock.rst
b/providers/common/ai/docs/connections/pydantic_ai_bedrock.rst
index 744400b2659..4815bf6d389 100644
--- a/providers/common/ai/docs/connections/pydantic_ai_bedrock.rst
+++ b/providers/common/ai/docs/connections/pydantic_ai_bedrock.rst
@@ -15,12 +15,12 @@
specific language governing permissions and limitations
under the License.
-.. _howto/connection:pydanticai-bedrock:
+.. _howto/connection:pydanticai_bedrock:
Pydantic AI (AWS Bedrock) Connection
=======================================
-The ``pydanticai-bedrock`` connection type configures access to
+The ``pydanticai_bedrock`` connection type configures access to
`AWS Bedrock <https://aws.amazon.com/bedrock/>`__ via the pydantic-ai
framework.
It backs ``PydanticAIBedrockHook``, the dedicated subclass of
``PydanticAIHook``
for Bedrock's AWS-style credentials — IAM keys, a bearer token, or the default
@@ -28,6 +28,29 @@ credential chain — none of which fit the plain ``api_key`` +
``base_url`` shap
that the generic :doc:`pydantic_ai` connection assumes. All fields live in
``extra``; the ``password`` and ``host`` fields are hidden in the connection
form.
+.. note::
+
+ This connection type was previously named ``pydanticai-bedrock``.
+
+ Connections stored as a URI or as JSON need no change: ``-`` is how ``_``
is
+ encoded in a URI scheme, so ``pydanticai-bedrock`` is decoded to
``pydanticai_bedrock``
+ on read and resolves as before. That covers ``AIRFLOW_CONN_*`` environment
+ variables and secrets backends such as HashiCorp Vault, AWS Secrets
Manager and
+ GCP Secret Manager.
+
+ A connection whose type is stored verbatim does need updating, because the
+ hyphen is preserved and no longer matches a registered hook. That means
rows in
+ the metadata database, including any created through the UI, and
connections
+ imported in object form from a local file:
+
+ .. code-block:: bash
+
+ airflow connections get <conn_id> -o json # confirm conn_type is
'pydanticai-bedrock'
+ airflow connections delete <conn_id>
+ airflow connections add <conn_id> --conn-type pydanticai_bedrock ...
+
+ In the UI, edit the connection and re-pick its type.
+
Default Connection IDs
----------------------
@@ -95,7 +118,7 @@ the instance role or environment:
.. code-block:: json
{
- "conn_type": "pydanticai-bedrock",
+ "conn_type": "pydanticai_bedrock",
"extra": "{\"model\": \"bedrock:us.anthropic.claude-opus-4-5\",
\"region_name\": \"us-east-1\"}"
}
@@ -104,7 +127,7 @@ the instance role or environment:
.. code-block:: json
{
- "conn_type": "pydanticai-bedrock",
+ "conn_type": "pydanticai_bedrock",
"extra": "{\"model\": \"bedrock:us.anthropic.claude-opus-4-5\",
\"region_name\": \"us-east-1\", \"aws_access_key_id\": \"AKIA...\",
\"aws_secret_access_key\": \"...\"}"
}
@@ -113,6 +136,6 @@ the instance role or environment:
.. code-block:: json
{
- "conn_type": "pydanticai-bedrock",
+ "conn_type": "pydanticai_bedrock",
"extra": "{\"model\": \"bedrock:us.anthropic.claude-opus-4-5\",
\"api_key\": \"<bearer-token>\"}"
}
diff --git a/providers/common/ai/docs/connections/pydantic_ai_vertex.rst
b/providers/common/ai/docs/connections/pydantic_ai_vertex.rst
index a242b202de7..5d91bb4ec7b 100644
--- a/providers/common/ai/docs/connections/pydantic_ai_vertex.rst
+++ b/providers/common/ai/docs/connections/pydantic_ai_vertex.rst
@@ -15,12 +15,12 @@
specific language governing permissions and limitations
under the License.
-.. _howto/connection:pydanticai-vertex:
+.. _howto/connection:pydanticai_vertex:
Pydantic AI (Google Vertex AI) Connection
============================================
-The ``pydanticai-vertex`` connection type configures access to
+The ``pydanticai_vertex`` connection type configures access to
`Google Vertex AI <https://cloud.google.com/vertex-ai>`__ via the pydantic-ai
framework. It backs ``PydanticAIVertexHook``, the dedicated subclass of
``PydanticAIHook`` for Google Cloud's project/location/service-account
@@ -29,6 +29,29 @@ shape that the generic :doc:`pydantic_ai` connection
assumes. All fields live
in ``extra``; the ``password`` and ``host`` fields are hidden in the connection
form.
+.. note::
+
+ This connection type was previously named ``pydanticai-vertex``.
+
+ Connections stored as a URI or as JSON need no change: ``-`` is how ``_``
is
+ encoded in a URI scheme, so ``pydanticai-vertex`` is decoded to
``pydanticai_vertex``
+ on read and resolves as before. That covers ``AIRFLOW_CONN_*`` environment
+ variables and secrets backends such as HashiCorp Vault, AWS Secrets
Manager and
+ GCP Secret Manager.
+
+ A connection whose type is stored verbatim does need updating, because the
+ hyphen is preserved and no longer matches a registered hook. That means
rows in
+ the metadata database, including any created through the UI, and
connections
+ imported in object form from a local file:
+
+ .. code-block:: bash
+
+ airflow connections get <conn_id> -o json # confirm conn_type is
'pydanticai-vertex'
+ airflow connections delete <conn_id>
+ airflow connections add <conn_id> --conn-type pydanticai_vertex ...
+
+ In the UI, edit the connection and re-pick its type.
+
Default Connection IDs
----------------------
@@ -114,7 +137,7 @@ environment:
.. code-block:: json
{
- "conn_type": "pydanticai-vertex",
+ "conn_type": "pydanticai_vertex",
"extra": "{\"model\": \"google-cloud:gemini-2.0-flash\", \"project\":
\"my-gcp-project\", \"location\": \"us-central1\"}"
}
@@ -123,6 +146,6 @@ environment:
.. code-block:: json
{
- "conn_type": "pydanticai-vertex",
+ "conn_type": "pydanticai_vertex",
"extra": "{\"model\": \"google-cloud:gemini-2.0-flash\", \"project\":
\"my-gcp-project\", \"location\": \"us-central1\", \"service_account_info\":
{\"type\": \"service_account\", \"project_id\": \"my-gcp-project\",
\"private_key\": \"<contents of the service account JSON key's private_key
field>\", \"client_email\": \"[email protected]\"}}"
}
diff --git a/providers/common/ai/docs/index.rst
b/providers/common/ai/docs/index.rst
index 117cfc8a910..fa9ea0e6c7f 100644
--- a/providers/common/ai/docs/index.rst
+++ b/providers/common/ai/docs/index.rst
@@ -48,7 +48,7 @@ When to use this provider
(OpenAI, Anthropic, Google, Bedrock, …) is picked by the connection
``llm_conn_id`` points
at — switching providers later is a connection change, not a Dag rewrite. Most
connections
use the generic ``pydanticai`` type, but Azure OpenAI, Bedrock, and Vertex AI
also have their
-own connection types (``pydanticai-azure``, ``pydanticai-bedrock``,
``pydanticai-vertex``) for
+own connection types (``pydanticai_azure``, ``pydanticai_bedrock``,
``pydanticai_vertex``) for
provider-specific authentication. Existing LangChain
tools aren't locked out either: pydantic-ai ships
``pydantic_ai.ext.langchain.LangChainToolset``
upstream, which wraps LangChain tools for a common.ai agent, and the
provider's own
diff --git a/providers/common/ai/provider.yaml
b/providers/common/ai/provider.yaml
index 1662b2f698d..bf4f9d0dbe2 100644
--- a/providers/common/ai/provider.yaml
+++ b/providers/common/ai/provider.yaml
@@ -182,7 +182,7 @@ connection-types:
- 'null'
- hook-class-name:
airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIAzureHook
hook-name: "Pydantic AI (Azure OpenAI)"
- connection-type: pydanticai-azure
+ connection-type: pydanticai_azure
external-services:
- Azure OpenAI
ui-field-behaviour:
@@ -213,7 +213,7 @@ connection-types:
- 'null'
- hook-class-name:
airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIBedrockHook
hook-name: "Pydantic AI (AWS Bedrock)"
- connection-type: pydanticai-bedrock
+ connection-type: pydanticai_bedrock
external-services:
- AWS Bedrock
ui-field-behaviour:
@@ -300,7 +300,7 @@ connection-types:
- 'null'
- hook-class-name:
airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIVertexHook
hook-name: "Pydantic AI (Google Vertex AI)"
- connection-type: pydanticai-vertex
+ connection-type: pydanticai_vertex
external-services:
- Google Vertex AI
ui-field-behaviour:
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py
b/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py
index 4aa7da1a98a..70fe47e607c 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py
@@ -155,7 +155,7 @@ def get_provider_info():
{
"hook-class-name":
"airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIAzureHook",
"hook-name": "Pydantic AI (Azure OpenAI)",
- "connection-type": "pydanticai-azure",
+ "connection-type": "pydanticai_azure",
"external-services": ["Azure OpenAI"],
"ui-field-behaviour": {
"hidden-fields": ["schema", "port", "login"],
@@ -181,7 +181,7 @@ def get_provider_info():
{
"hook-class-name":
"airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIBedrockHook",
"hook-name": "Pydantic AI (AWS Bedrock)",
- "connection-type": "pydanticai-bedrock",
+ "connection-type": "pydanticai_bedrock",
"external-services": ["AWS Bedrock"],
"ui-field-behaviour": {
"hidden-fields": ["schema", "port", "login", "host",
"password"],
@@ -246,7 +246,7 @@ def get_provider_info():
{
"hook-class-name":
"airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIVertexHook",
"hook-name": "Pydantic AI (Google Vertex AI)",
- "connection-type": "pydanticai-vertex",
+ "connection-type": "pydanticai_vertex",
"external-services": ["Google Vertex AI"],
"ui-field-behaviour": {
"hidden-fields": ["schema", "port", "login", "host",
"password"],
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
b/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
index 4b3ce7357c5..6b5658c42e5 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
@@ -317,7 +317,7 @@ class PydanticAIAzureHook(PydanticAIHook):
:param model_id: Model identifier, e.g. ``"azure:gpt-4o"``.
"""
- conn_type = "pydanticai-azure"
+ conn_type = "pydanticai_azure"
default_conn_name = "pydanticai_azure_default"
hook_name = "Pydantic AI (Azure OpenAI)"
@@ -385,7 +385,7 @@ class PydanticAIBedrockHook(PydanticAIHook):
:param model_id: Model identifier, e.g.
``"bedrock:us.anthropic.claude-opus-4-5"``.
"""
- conn_type = "pydanticai-bedrock"
+ conn_type = "pydanticai_bedrock"
default_conn_name = "pydanticai_bedrock_default"
hook_name = "Pydantic AI (AWS Bedrock)"
@@ -477,7 +477,7 @@ class PydanticAIVertexHook(PydanticAIHook):
:param model_id: Model identifier, e.g.
``"google-cloud:gemini-2.0-flash"``.
"""
- conn_type = "pydanticai-vertex"
+ conn_type = "pydanticai_vertex"
default_conn_name = "pydanticai_vertex_default"
hook_name = "Pydantic AI (Google Vertex AI)"
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
index 36cd616c1b3..a2f68bd730b 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
@@ -147,7 +147,7 @@ class LLMOperator(BaseOperator, LLMApprovalMixin):
Delegates to :meth:`~PydanticAIHook.get_hook` which looks up
the connection's ``conn_type`` and instantiates the matching subclass
(e.g.
:class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIAzureHook`
- for ``pydanticai-azure`` connections).
+ for ``pydanticai_azure`` connections).
"""
hook_params = {
"model_id": self.model_id,
diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
b/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
index af13f3ceb22..76cf0e0dd03 100644
--- a/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
+++ b/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
@@ -520,7 +520,7 @@ class TestPydanticAIAzureHook:
"""Tests for PydanticAIAzureHook."""
def test_conn_type(self):
- assert PydanticAIAzureHook.conn_type == "pydanticai-azure"
+ assert PydanticAIAzureHook.conn_type == "pydanticai_azure"
def test_hook_name(self):
assert "Azure" in PydanticAIAzureHook.hook_name
@@ -581,7 +581,7 @@ class TestPydanticAIAzureHook:
hook = PydanticAIAzureHook(llm_conn_id="azure_test")
conn = Connection(
conn_id="azure_test",
- conn_type="pydanticai-azure",
+ conn_type="pydanticai_azure",
password="azure-key",
host="https://myresource.openai.azure.com";,
extra=json.dumps({"model": "azure:gpt-4o", "api_version":
"2024-07-01-preview"}),
@@ -604,7 +604,7 @@ class TestPydanticAIAzureHook:
hook = PydanticAIAzureHook(llm_conn_id="azure_test")
conn = Connection(
conn_id="azure_test",
- conn_type="pydanticai-azure",
+ conn_type="pydanticai_azure",
extra=json.dumps({"model": "azure:gpt-4o"}),
)
with patch.object(hook, "get_connection", return_value=conn):
@@ -617,7 +617,7 @@ class TestPydanticAIBedrockHook:
"""Tests for PydanticAIBedrockHook."""
def test_conn_type(self):
- assert PydanticAIBedrockHook.conn_type == "pydanticai-bedrock"
+ assert PydanticAIBedrockHook.conn_type == "pydanticai_bedrock"
def test_hook_name(self):
assert "Bedrock" in PydanticAIBedrockHook.hook_name
@@ -657,7 +657,7 @@ class TestPydanticAIBedrockHook:
hook = PydanticAIBedrockHook(llm_conn_id="bedrock_test")
conn = Connection(
conn_id="bedrock_test",
- conn_type="pydanticai-bedrock",
+ conn_type="pydanticai_bedrock",
extra=json.dumps({"model":
"bedrock:us.anthropic.claude-opus-4-5"}),
)
with patch.object(hook, "get_connection", return_value=conn):
@@ -675,7 +675,7 @@ class TestPydanticAIBedrockHook:
hook = PydanticAIBedrockHook(llm_conn_id="bedrock_test")
conn = Connection(
conn_id="bedrock_test",
- conn_type="pydanticai-bedrock",
+ conn_type="pydanticai_bedrock",
extra=json.dumps(
{
"model": "bedrock:us.anthropic.claude-opus-4-5",
@@ -747,7 +747,7 @@ class TestPydanticAIVertexHook:
"""Tests for PydanticAIVertexHook."""
def test_conn_type(self):
- assert PydanticAIVertexHook.conn_type == "pydanticai-vertex"
+ assert PydanticAIVertexHook.conn_type == "pydanticai_vertex"
def test_hook_name(self):
assert "Vertex" in PydanticAIVertexHook.hook_name
@@ -858,7 +858,7 @@ class TestPydanticAIVertexHook:
hook = PydanticAIVertexHook(llm_conn_id="vertex_test")
conn = Connection(
conn_id="vertex_test",
- conn_type="pydanticai-vertex",
+ conn_type="pydanticai_vertex",
extra=json.dumps({"model": "google-cloud:gemini-2.0-flash"}),
)
with patch.object(hook, "get_connection", return_value=conn):
@@ -876,7 +876,7 @@ class TestPydanticAIVertexHook:
hook = PydanticAIVertexHook(llm_conn_id="vertex_test")
conn = Connection(
conn_id="vertex_test",
- conn_type="pydanticai-vertex",
+ conn_type="pydanticai_vertex",
extra=json.dumps(
{
"model": "google-cloud:gemini-2.0-flash",
@@ -934,7 +934,7 @@ class TestPydanticAIVertexHook:
hook = PydanticAIVertexHook(llm_conn_id="vertex_test")
conn = Connection(
conn_id="vertex_test",
- conn_type="pydanticai-vertex",
+ conn_type="pydanticai_vertex",
extra=json.dumps(
{
"model": "google-cloud:gemini-2.0-flash",
@@ -976,7 +976,7 @@ class TestPydanticAIVertexHook:
"""
connection_types = get_provider_info()["connection-types"]
vertex_conn_fields = next(
- c["conn-fields"] for c in connection_types if c["connection-type"]
== "pydanticai-vertex"
+ c["conn-fields"] for c in connection_types if c["connection-type"]
== "pydanticai_vertex"
)
description = vertex_conn_fields["model"]["description"]
prefix = _extract_google_cloud_prefix(description)
@@ -994,7 +994,7 @@ class TestPydanticAIVertexHook:
"""
connection_types = get_provider_info()["connection-types"]
vertex_connection_type = next(
- c for c in connection_types if c["connection-type"] ==
"pydanticai-vertex"
+ c for c in connection_types if c["connection-type"] ==
"pydanticai_vertex"
)
placeholder =
vertex_connection_type["ui-field-behaviour"]["placeholders"]["extra"]
prefix = _extract_google_cloud_prefix(placeholder)
@@ -1011,3 +1011,70 @@ class TestPydanticAIVertexHook:
placeholder =
PydanticAIVertexHook.get_ui_field_behaviour()["placeholders"]["extra"]
prefix = _extract_google_cloud_prefix(placeholder)
_assert_prefix_is_known_provider(prefix)
+
+
+ALL_HOOKS = [PydanticAIHook, PydanticAIAzureHook, PydanticAIBedrockHook,
PydanticAIVertexHook]
+DECLARED_CONNECTION_TYPES = [c["connection-type"] for c in
get_provider_info()["connection-types"]]
+
+
+class TestConnTypeResolution:
+ """
+ Every hook is registered under the literal ``connection-type`` string from
+ ``provider.yaml``, but ``Connection.from_uri`` and
``Connection.from_json`` rewrite
+ ``-`` to ``_`` before the lookup happens. A hyphen in ``conn_type``
therefore makes
+ the hook unreachable from every secrets backend, which is what happened to
the three
+ vendor types (apache/airflow#72316): only a connection read straight out
of the
+ metadata DB kept its hyphen and resolved.
+ """
+
+ def test_provider_declares_connection_types(self):
+ """Keeps the round-trip guard below from passing vacuously."""
+ assert DECLARED_CONNECTION_TYPES
+
+ @pytest.mark.parametrize("conn_type", DECLARED_CONNECTION_TYPES)
+ def test_declared_connection_type_survives_a_uri_round_trip(self,
conn_type):
+ """Guards every connection-type this provider declares, current and
future."""
+ source = Connection(conn_id="c", conn_type=conn_type)
+
+ parsed = Connection(conn_id="c", uri=source.get_uri())
+
+ assert parsed.conn_type == conn_type, (
+ f"connection-type {conn_type!r} does not survive URI
serialization, so its hook "
+ "cannot be looked up from any secrets backend; declare it with
underscores"
+ )
+
+ @pytest.mark.parametrize("hook_class", ALL_HOOKS, ids=lambda c: c.__name__)
+ def test_hook_resolves_from_uri(self, hook_class):
+ conn = Connection(conn_id="c", conn_type=hook_class.conn_type,
host="example.com")
+
+ round_tripped = Connection(conn_id="c", uri=conn.get_uri())
+
+ assert round_tripped.conn_type == hook_class.conn_type
+ assert type(round_tripped.get_hook()) is hook_class
+
+ @pytest.mark.parametrize("hook_class", ALL_HOOKS, ids=lambda c: c.__name__)
+ def test_hook_resolves_from_json(self, hook_class):
+ conn = Connection(conn_id="c", conn_type=hook_class.conn_type,
host="example.com")
+
+ round_tripped = Connection.from_json(conn.as_json(), conn_id="c")
+
+ assert round_tripped.conn_type == hook_class.conn_type
+ assert type(round_tripped.get_hook()) is hook_class
+
+ @pytest.mark.parametrize("hook_class", ALL_HOOKS, ids=lambda c: c.__name__)
+ @pytest.mark.parametrize("serializer", ["uri", "json"])
+ def test_hook_resolves_from_environment_variable(self, hook_class,
serializer, monkeypatch):
+ """
+ ``AIRFLOW_CONN_*`` is enabled by default and is how most deployments
define
+ connections, so it is the widest blast radius for a conn_type that
does not
+ round-trip.
+ """
+ conn_id = f"vendor_{hook_class.conn_type}_{serializer}"
+ conn = Connection(conn_id=conn_id, conn_type=hook_class.conn_type,
host="example.com")
+ serialized = conn.get_uri() if serializer == "uri" else conn.as_json()
+ monkeypatch.setenv(f"AIRFLOW_CONN_{conn_id.upper()}", serialized)
+
+ resolved = Connection.get_connection_from_secrets(conn_id)
+
+ assert resolved.conn_type == hook_class.conn_type
+ assert type(resolved.get_hook()) is hook_class
