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 <uranusjr@gmail.com>

Author: amoghrajesh

task-sdk/src/airflow/sdk/definitions/connection.py

@@ -20,7 +20,7 @@
 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
 
 
-@attrs.define
+@attrs.define(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

task-sdk/tests/task_sdk/definitions/test_connection.py

@@ -377,6 +377,20 @@ def test_from_uri_invalid_protocol_host_error(self):
         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&param2=value2"