diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py index 34cfe47334893..ebddbcaa85187 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/connections.py @@ -29,6 +29,8 @@ from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel, make_partial_model from airflow.configuration import conf +PortNumber = Annotated[int, Field(ge=0, le=65535)] + # Response Models class ConnectionResponse(BaseModel): @@ -191,7 +193,7 @@ class ConnectionBody(StrictBaseModel): host: str | None = Field(default=None) login: str | None = Field(default=None) schema_: str | None = Field(None, alias="schema") - port: int | None = Field(default=None) + port: PortNumber | None = Field(default=None) password: str | None = Field(default=None) extra: str | None = Field(default=None) team_name: str | None = Field(max_length=50, default=None) diff --git a/airflow-core/src/airflow/models/connection.py b/airflow-core/src/airflow/models/connection.py index 1b4b0f8f86768..cd294a9f88fc9 100644 --- a/airflow-core/src/airflow/models/connection.py +++ b/airflow-core/src/airflow/models/connection.py @@ -28,7 +28,7 @@ from urllib.parse import parse_qsl, quote, unquote, urlencode, urlsplit from sqlalchemy import ForeignKey, Integer, String, Text, select -from sqlalchemy.orm import Mapped, mapped_column, reconstructor +from sqlalchemy.orm import Mapped, mapped_column, reconstructor, validates from airflow._shared.module_loading import import_string from airflow._shared.secrets_backend.base import call_secrets_backend_method @@ -60,6 +60,8 @@ class AirflowSecretsBackendAccessDenied(PermissionError): # type: ignore[no-red # # You can try the regex here: https://regex101.com/r/69033B/1 RE_SANITIZE_CONN_ID = re.compile(r"^[\w#!()\-.:/\\]{1,}$") +PORT_MIN = 0 +PORT_MAX = 65535 # the conn ID max len should be 250 CONN_ID_MAX_LEN: int = 250 @@ -219,6 +221,33 @@ def _validate_extra(extra, conn_id) -> None: raise ValueError(f"Encountered non-JSON in `extra` field for connection {conn_id!r}.") return None + @staticmethod + def _validate_port(port: int | None) -> int | None: + if port is None: + return None + if type(port) is not int: + raise ValueError(f"Expected integer value for `port`, but got {port!r} instead.") + if not PORT_MIN <= port <= PORT_MAX: + raise ValueError( + f"Expected value for `port` in the range {PORT_MIN}-{PORT_MAX}, but got {port!r} instead." + ) + return port + + @classmethod + def _coerce_port(cls, port) -> int | None: + if port is None or port == "": + return None + if isinstance(port, str): + try: + port = int(port) + except ValueError: + raise ValueError(f"Expected integer value for `port`, but got {port!r} instead.") from None + return cls._validate_port(port) + + @validates("port") + def _validate_port_assignment(self, _, port: int | None) -> int | None: + return self._validate_port(port) + @reconstructor def on_db_load(self): if self.password: @@ -335,7 +364,7 @@ def get_uri(self) -> str: if host_to_use: host_block += quote(host_to_use, safe="") - if self.port: + if self.port is not None: if host_block == "" and authority_block == "": host_block += f"@:{self.port}" else: @@ -586,11 +615,8 @@ def from_json(cls, value, conn_id=None) -> Connection: if conn_type: kwargs["conn_type"] = cls._normalize_conn_type(conn_type) port = kwargs.pop("port", None) - if port: - try: - kwargs["port"] = int(port) - except ValueError: - raise ValueError(f"Expected integer value for `port`, but got {port!r} instead.") + if port is not None: + kwargs["port"] = cls._coerce_port(port) return Connection(conn_id=conn_id, **kwargs) def as_json(self) -> str: diff --git a/airflow-core/tests/unit/always/test_connection.py b/airflow-core/tests/unit/always/test_connection.py index 8b56791d7c5fe..3267861c4459c 100644 --- a/airflow-core/tests/unit/always/test_connection.py +++ b/airflow-core/tests/unit/always/test_connection.py @@ -620,8 +620,11 @@ def test_from_json_conn_type(self, val, expected): @pytest.mark.parametrize( ("val", "expected"), [ + ('{"port": 0}', 0), ('{"port": 1}', 1), ('{"port": "1"}', 1), + ('{"port": "0"}', 0), + ('{"port": ""}', None), ('{"port": null}', None), ], ) @@ -629,6 +632,11 @@ def test_from_json_port(self, val, expected): """Two conn_type normalizations are applied: replace - with _ and postgresql with postgres""" assert Connection.from_json(val).port == expected + @pytest.mark.parametrize("val", ['{"port": -1}', '{"port": 65536}', '{"port": 1.5}', '{"port": "1.5"}']) + def test_from_json_rejects_invalid_port(self, val): + with pytest.raises(ValueError, match="port"): + Connection.from_json(val) + @pytest.mark.parametrize( ("val", "expected"), [ diff --git a/airflow-core/tests/unit/api_fastapi/core_api/datamodels/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/datamodels/test_connections.py index 072f35167fd29..5307cf6198d79 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/datamodels/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/datamodels/test_connections.py @@ -19,7 +19,7 @@ import pytest from pydantic import ValidationError -from airflow.api_fastapi.core_api.datamodels.connections import ConnectionResponse +from airflow.api_fastapi.core_api.datamodels.connections import ConnectionBody, ConnectionResponse def _payload(**overrides): @@ -56,3 +56,15 @@ def test_redact_extra_rejects_non_json(extra): """ with pytest.raises(ValidationError): ConnectionResponse.model_validate(_payload(extra=extra)) + + +@pytest.mark.parametrize("port", [0, 65535, None]) +def test_connection_body_allows_valid_port_boundaries(port): + body = ConnectionBody(connection_id="test_conn", conn_type="http", port=port) + assert body.port == port + + +@pytest.mark.parametrize("port", [-1, 65536]) +def test_connection_body_rejects_invalid_port_boundaries(port): + with pytest.raises(ValidationError): + ConnectionBody(connection_id="test_conn", conn_type="http", port=port) diff --git a/airflow-core/tests/unit/models/test_connection.py b/airflow-core/tests/unit/models/test_connection.py index 94cabe5e4daf4..f513f944f515b 100644 --- a/airflow-core/tests/unit/models/test_connection.py +++ b/airflow-core/tests/unit/models/test_connection.py @@ -320,6 +320,16 @@ def test_get_uri_conn_type_warning(self, connection: Connection, expected_warned f"RFC3986 warning not expected for connection '{connection.conn_id}'." ) + @pytest.mark.parametrize("port", [0, 65535, None]) + def test_allows_valid_port_boundaries(self, port): + conn = Connection(conn_id="test-port", conn_type="http", port=port) + assert conn.port == port + + @pytest.mark.parametrize("port", [-1, 65536, "123", 1.5, ""]) + def test_rejects_invalid_direct_port_values(self, port): + with pytest.raises(ValueError, match="port"): + Connection(conn_id="test-port", conn_type="http", port=port) + @pytest.mark.parametrize( ("connection", "expected_conn_id"), [ diff --git a/task-sdk/src/airflow/sdk/definitions/connection.py b/task-sdk/src/airflow/sdk/definitions/connection.py index 06a95a3b868b8..d4c5324df2842 100644 --- a/task-sdk/src/airflow/sdk/definitions/connection.py +++ b/task-sdk/src/airflow/sdk/definitions/connection.py @@ -29,6 +29,19 @@ from airflow.sdk.providers_manager_runtime import ProvidersManagerTaskRuntime log = logging.getLogger(__name__) +PORT_MIN = 0 +PORT_MAX = 65535 + + +def _validate_port(_, __, value: int | None) -> None: + if value is None: + return + if type(value) is not int: + raise ValueError(f"Expected integer value for `port`, but got {value!r} instead.") + if not PORT_MIN <= value <= PORT_MAX: + raise ValueError( + f"Expected value for `port` in the range {PORT_MIN}-{PORT_MAX}, but got {value!r} instead." + ) def _parse_netloc_to_hostname(uri_parts): @@ -118,7 +131,7 @@ class Connection: schema: str | None = None login: str | None = None password: str | None = None - port: int | None = None + port: int | None = attrs.field(default=None, validator=attrs.validators.optional(_validate_port)) extra: str | None = None EXTRA_KEY = "__extra__" @@ -153,6 +166,18 @@ def __init__(self, *, conn_id: str, uri: str | None = None, **kwargs) -> None: else: self.__dict__.update(attrs.asdict(self.from_uri(uri, conn_id=conn_id), recurse=False)) + @classmethod + def _coerce_port(cls, port) -> int | None: + if port is None or port == "": + return None + if isinstance(port, str): + try: + port = int(port) + except ValueError: + raise ValueError(f"Expected integer value for `port`, but got {port!r} instead.") from None + _validate_port(None, None, port) + return port + def get_uri(self) -> str: """Generate and return connection in URI format.""" from urllib.parse import parse_qsl @@ -191,7 +216,7 @@ def get_uri(self) -> str: host_block = "" if host_to_use: host_block += quote(host_to_use, safe="") - if self.port: + if self.port is not None: if host_block == "" and authority_block == "": host_block += f"@:{self.port}" else: @@ -355,11 +380,8 @@ def from_json(cls, value, conn_id=None) -> Connection: if conn_type: kwargs["conn_type"] = cls._normalize_conn_type(conn_type) port = kwargs.pop("port", None) - if port: - try: - kwargs["port"] = int(port) - except ValueError: - raise ValueError(f"Expected integer value for `port`, but got {port!r} instead.") + if port is not None: + kwargs["port"] = cls._coerce_port(port) return cls(conn_id=conn_id, **kwargs) def as_json(self) -> str: diff --git a/task-sdk/tests/task_sdk/definitions/test_connection.py b/task-sdk/tests/task_sdk/definitions/test_connection.py index c973fc58c91a4..3cbe35cc86d2b 100644 --- a/task-sdk/tests/task_sdk/definitions/test_connection.py +++ b/task-sdk/tests/task_sdk/definitions/test_connection.py @@ -198,6 +198,34 @@ def test_from_json(self): assert connection.host == "localhost" assert connection.port == 5432 + @pytest.mark.parametrize( + ("payload", "expected_port"), + [ + ({"port": 0}, 0), + ({"port": "0"}, 0), + ({"port": 65535}, 65535), + ({"port": ""}, None), + ], + ) + def test_from_json_port_boundaries(self, payload, expected_port): + connection = Connection.from_json(json.dumps(payload), conn_id="test_conn") + assert connection.port == expected_port + + @pytest.mark.parametrize("payload", [{"port": -1}, {"port": 65536}, {"port": 1.5}, {"port": "1.5"}]) + def test_from_json_rejects_invalid_port(self, payload): + with pytest.raises(ValueError, match="port"): + Connection.from_json(json.dumps(payload), conn_id="test_conn") + + @pytest.mark.parametrize("port", [0, 65535, None]) + def test_direct_constructor_allows_valid_port_boundaries(self, port): + connection = Connection(conn_id="test_conn", conn_type="http", port=port) + assert connection.port == port + + @pytest.mark.parametrize("port", [-1, 65536, "123", 1.5, ""]) + def test_direct_constructor_rejects_invalid_port_values(self, port): + with pytest.raises(ValueError, match="port"): + Connection(conn_id="test_conn", conn_type="http", port=port) + def test_from_json_without_conn_type(self): """Test that from_json works without conn_type (backward compatibility with AF 2).""" json_data = {