This is an automated email from the ASF dual-hosted git repository.
amoghdesai 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 8a3a5222e39 Add support to create connections using uri in SDK (#62211)
8a3a5222e39 is described below
commit 8a3a5222e39d83c842b5e44ffe8475c0309380ea
Author: Amogh Desai <[email protected]>
AuthorDate: Fri Feb 27 19:03:44 2026 +0530
Add support to create connections using uri in SDK (#62211)
* Add support to create connections using uri in SDK
* Add support to create connections using uri in SDK
* Add support to create connections using uri in SDK
* fixing unit tests
* Allow URI without adding it as a field
* fixing failing tests and mypy
* remove unwanted line
---------
Co-authored-by: Tzu-ping Chung <[email protected]>
---
task-sdk/src/airflow/sdk/definitions/connection.py | 35 ++++++++++++++++++++--
.../tests/task_sdk/definitions/test_connection.py | 14 +++++++++
2 files changed, 47 insertions(+), 2 deletions(-)
diff --git a/task-sdk/src/airflow/sdk/definitions/connection.py
b/task-sdk/src/airflow/sdk/definitions/connection.py
index 55fc269010a..21eed4216f5 100644
--- a/task-sdk/src/airflow/sdk/definitions/connection.py
+++ b/task-sdk/src/airflow/sdk/definitions/connection.py
@@ -20,7 +20,7 @@ from __future__ import annotations
import json
import logging
from json import JSONDecodeError
-from typing import Any
+from typing import Any, overload
from urllib.parse import parse_qsl, quote, unquote, urlencode, urlsplit
import attrs
@@ -93,7 +93,7 @@ def _prune_dict(val: Any, mode="strict"):
return val
[email protected]
[email protected](slots=False)
class Connection:
"""
A connection to an external data source.
@@ -108,6 +108,7 @@ class Connection:
:param port: The port number.
:param extra: Extra metadata. Non-standard data such as private/SSH keys
can be saved here. JSON
encoded object.
+ :param uri: URI address describing connection parameters.
"""
conn_id: str
@@ -122,6 +123,36 @@ class Connection:
EXTRA_KEY = "__extra__"
+ @overload
+ def __init__(self, *, conn_id: str, uri: str) -> None: ...
+
+ @overload
+ def __init__(
+ self,
+ *,
+ conn_id: str,
+ conn_type: str | None = None,
+ description: str | None = None,
+ host: str | None = None,
+ schema: str | None = None,
+ login: str | None = None,
+ password: str | None = None,
+ port: int | None = None,
+ extra: str | None = None,
+ ) -> None: ...
+
+ def __init__(self, *, conn_id: str, uri: str | None = None, **kwargs) ->
None:
+ if uri is not None and kwargs:
+ raise AirflowException(
+ "You must create an object using the URI or individual values "
+ "(conn_type, host, login, password, schema, port or extra). "
+ "You can't mix these two ways to create this object."
+ )
+ if uri is None:
+ self.__attrs_init__(conn_id=conn_id, **kwargs) # type:
ignore[attr-defined]
+ else:
+ self.__dict__.update(self.from_uri(uri,
conn_id=conn_id).to_dict(validate=False))
+
def get_uri(self) -> str:
"""Generate and return connection in URI format."""
from urllib.parse import parse_qsl
diff --git a/task-sdk/tests/task_sdk/definitions/test_connection.py
b/task-sdk/tests/task_sdk/definitions/test_connection.py
index cb2455a1f7c..470546d57da 100644
--- a/task-sdk/tests/task_sdk/definitions/test_connection.py
+++ b/task-sdk/tests/task_sdk/definitions/test_connection.py
@@ -377,6 +377,20 @@ class TestConnectionFromUri:
with pytest.raises(AirflowException, match="Invalid connection
string"):
Connection.from_uri(uri, conn_id="test_conn")
+ def test_connection_constructor_with_uri(self):
+ """Test Connection(uri=..., conn_id=...) constructor form."""
+ conn = Connection(conn_id="test_conn",
uri="postgres://user:pass@host:5432/db")
+
+ assert conn.conn_id == "test_conn"
+ assert conn.conn_type == "postgres"
+ assert conn.host == "host"
+ assert conn.login == "user"
+ assert conn.password == "pass"
+ assert conn.port == 5432
+ assert conn.schema == "db"
+ # uri should not exist as an attribute (it's init-only)
+ assert not hasattr(conn, "uri")
+
def test_from_uri_roundtrip(self):
"""Test that from_uri and get_uri are inverse operations."""
original_uri =
"postgres://user:pass@host:5432/db?param1=value1¶m2=value2"
