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 4bbff9069bb36..ee22c09d116c4 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 @@ -41,6 +41,7 @@ class ConnectionResponse(BaseModel): port: int | None password: str | None extra: str | None + team_name: str | None @field_validator("password", mode="after") @classmethod @@ -136,6 +137,7 @@ class ConnectionBody(StrictBaseModel): port: int | 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) @field_validator("extra") @classmethod diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index 5bcc026a26e56..36002f11c3490 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -9732,6 +9732,12 @@ components: - type: string - type: 'null' title: Extra + team_name: + anyOf: + - type: string + maxLength: 50 + - type: 'null' + title: Team Name additionalProperties: false type: object required: @@ -9798,6 +9804,11 @@ components: - type: string - type: 'null' title: Extra + team_name: + anyOf: + - type: string + - type: 'null' + title: Team Name type: object required: - connection_id @@ -9809,6 +9820,7 @@ components: - port - password - extra + - team_name title: ConnectionResponse description: Connection serializer for responses. ConnectionTestResponse: diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index 419f3b2da63cc..d4a506eab4255 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -1545,6 +1545,18 @@ export const $ConnectionBody = { } ], title: 'Extra' + }, + team_name: { + anyOf: [ + { + type: 'string', + maxLength: 50 + }, + { + type: 'null' + } + ], + title: 'Team Name' } }, additionalProperties: false, @@ -1660,10 +1672,21 @@ export const $ConnectionResponse = { } ], title: 'Extra' + }, + team_name: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Team Name' } }, type: 'object', - required: ['connection_id', 'conn_type', 'description', 'host', 'login', 'schema', 'port', 'password', 'extra'], + required: ['connection_id', 'conn_type', 'description', 'host', 'login', 'schema', 'port', 'password', 'extra', 'team_name'], title: 'ConnectionResponse', description: 'Connection serializer for responses.' } as const; diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index 05391098ac158..fddddb9bc8ae8 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -468,6 +468,7 @@ export type ConnectionBody = { port?: number | null; password?: string | null; extra?: string | null; + team_name?: string | null; }; /** @@ -491,6 +492,7 @@ export type ConnectionResponse = { port: number | null; password: string | null; extra: string | null; + team_name: string | null; }; /** diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index 3ebeecb677d9a..3a6a9f8b1c0c3 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -21,10 +21,11 @@ from unittest import mock import pytest +from sqlalchemy.orm import Session from airflow.models import Connection from airflow.secrets.environment_variables import CONN_ENV_PREFIX -from airflow.utils.session import provide_session +from airflow.utils.session import NEW_SESSION, provide_session from tests_common.test_utils.api_fastapi import _check_last_log from tests_common.test_utils.asserts import assert_queries_count @@ -56,7 +57,7 @@ @provide_session -def _create_connection(session) -> None: +def _create_connection(team_name: str | None = None, session: Session = NEW_SESSION) -> None: connection_model = Connection( conn_id=TEST_CONN_ID, conn_type=TEST_CONN_TYPE, @@ -64,13 +65,14 @@ def _create_connection(session) -> None: host=TEST_CONN_HOST, port=TEST_CONN_PORT, login=TEST_CONN_LOGIN, + team_name=team_name, ) session.add(connection_model) @provide_session -def _create_connections(session) -> None: - _create_connection(session) +def _create_connections(session: Session = NEW_SESSION) -> None: + _create_connection(session=session) connection_model_2 = Connection( conn_id=TEST_CONN_ID_2, conn_type=TEST_CONN_TYPE_2, @@ -92,8 +94,8 @@ def setup(self) -> None: def teardown_method(self) -> None: clear_db_connections() - def create_connection(self): - _create_connection() + def create_connection(self, team_name: str | None = None): + _create_connection(team_name=team_name) def create_connections(self): _create_connections() @@ -126,13 +128,14 @@ def test_delete_should_respond_404(self, test_client): class TestGetConnection(TestConnectionEndpoint): - def test_get_should_respond_200(self, test_client, session): - self.create_connection() + def test_get_should_respond_200(self, test_client, testing_team, session): + self.create_connection(team_name=testing_team.name) response = test_client.get(f"/connections/{TEST_CONN_ID}") assert response.status_code == 200 body = response.json() assert body["connection_id"] == TEST_CONN_ID assert body["conn_type"] == TEST_CONN_TYPE + assert body["team_name"] == testing_team.name def test_should_respond_401(self, unauthenticated_test_client): response = unauthenticated_test_client.get(f"/connections/{TEST_CONN_ID}") @@ -278,6 +281,25 @@ def test_post_should_respond_201(self, test_client, session, body): assert len(connection) == 1 _check_last_log(session, dag_id=None, event="post_connection", logical_date=None) + def test_post_should_respond_201_with_team(self, test_client, session, testing_team): + response = test_client.post( + "/connections", + json={"connection_id": TEST_CONN_ID, "conn_type": TEST_CONN_TYPE, "team_name": testing_team.name}, + ) + assert response.status_code == 201 + assert response.json() == { + "connection_id": TEST_CONN_ID, + "conn_type": TEST_CONN_TYPE, + "description": None, + "extra": None, + "host": None, + "login": None, + "password": None, + "port": None, + "schema": None, + "team_name": testing_team.name, + } + def test_should_respond_401(self, unauthenticated_test_client): response = unauthenticated_test_client.post("/connections", json={}) assert response.status_code == 401 @@ -343,6 +365,7 @@ def test_post_should_respond_already_exist(self, test_client, body): "password": "***", "port": None, "schema": None, + "team_name": None, }, ), ( @@ -357,6 +380,7 @@ def test_post_should_respond_already_exist(self, test_client, body): "password": "***", "port": None, "schema": None, + "team_name": None, }, ), ( @@ -376,6 +400,7 @@ def test_post_should_respond_already_exist(self, test_client, body): "password": "***", "port": None, "schema": None, + "team_name": None, }, ), ], @@ -403,6 +428,7 @@ class TestPatchConnection(TestConnectionEndpoint): "password": None, "port": TEST_CONN_PORT, "schema": None, + "team_name": None, }, ), ( @@ -417,6 +443,7 @@ class TestPatchConnection(TestConnectionEndpoint): "password": None, "port": TEST_CONN_PORT, "schema": None, + "team_name": None, }, ), ( @@ -436,6 +463,7 @@ class TestPatchConnection(TestConnectionEndpoint): "password": None, "port": 80, "schema": None, + "team_name": None, }, ), ( @@ -450,6 +478,7 @@ class TestPatchConnection(TestConnectionEndpoint): "password": None, "port": TEST_CONN_PORT, "schema": None, + "team_name": None, }, ), ( @@ -464,6 +493,7 @@ class TestPatchConnection(TestConnectionEndpoint): "password": None, "port": 80, "schema": None, + "team_name": None, }, ), ( @@ -484,6 +514,7 @@ class TestPatchConnection(TestConnectionEndpoint): "password": "test_password_patch", "port": 80, "schema": None, + "team_name": None, }, ), ( @@ -505,6 +536,7 @@ class TestPatchConnection(TestConnectionEndpoint): "password": None, "port": 80, "schema": None, + "team_name": None, }, ), ( @@ -524,6 +556,7 @@ class TestPatchConnection(TestConnectionEndpoint): "password": None, "port": TEST_CONN_PORT, "schema": "http_patch", + "team_name": None, }, ), ( @@ -548,11 +581,11 @@ class TestPatchConnection(TestConnectionEndpoint): "password": None, "port": None, "schema": None, + "team_name": None, }, ), ], ) - @provide_session def test_patch_should_respond_200( self, test_client, body: dict[str, str], expected_result: dict[str, str], session ): @@ -564,6 +597,29 @@ def test_patch_should_respond_200( assert response.json() == expected_result + def test_patch_with_team_should_respond_200(self, test_client, testing_team, session): + self.create_connection() + + response = test_client.patch( + f"/connections/{TEST_CONN_ID}", + json={"connection_id": TEST_CONN_ID, "conn_type": "new_type", "team_name": testing_team.name}, + ) + assert response.status_code == 200 + _check_last_log(session, dag_id=None, event="patch_connection", logical_date=None) + + assert response.json() == { + "conn_type": "new_type", + "connection_id": TEST_CONN_ID, + "description": TEST_CONN_DESCRIPTION, + "extra": None, + "host": TEST_CONN_HOST, + "login": TEST_CONN_LOGIN, + "password": None, + "port": TEST_CONN_PORT, + "schema": None, + "team_name": testing_team.name, + } + def test_should_respond_401(self, unauthenticated_test_client): response = unauthenticated_test_client.patch(f"/connections/{TEST_CONN_ID}", json={}) assert response.status_code == 401 @@ -593,6 +649,7 @@ def test_should_respond_403(self, unauthorized_test_client): "schema": None, "password": None, "description": TEST_CONN_DESCRIPTION, + "team_name": None, }, {"update_mask": ["login", "port"]}, ), @@ -614,6 +671,7 @@ def test_should_respond_403(self, unauthorized_test_client): "schema": None, "password": None, "description": TEST_CONN_DESCRIPTION, + "team_name": None, }, {"update_mask": ["login", "port"]}, ), @@ -629,6 +687,7 @@ def test_should_respond_403(self, unauthorized_test_client): "schema": None, "password": None, "description": TEST_CONN_DESCRIPTION, + "team_name": None, }, {"update_mask": ["host"]}, ), @@ -649,6 +708,7 @@ def test_should_respond_403(self, unauthorized_test_client): "schema": None, "password": None, "description": TEST_CONN_DESCRIPTION, + "team_name": None, }, {"update_mask": ["host", "port"]}, ), @@ -664,6 +724,7 @@ def test_should_respond_403(self, unauthorized_test_client): "schema": None, "password": None, "description": TEST_CONN_DESCRIPTION, + "team_name": None, }, {"update_mask": ["login"]}, ), @@ -684,6 +745,7 @@ def test_should_respond_403(self, unauthorized_test_client): "password": None, "schema": None, "description": TEST_CONN_DESCRIPTION, + "team_name": None, }, {"update_mask": ["host"]}, ), @@ -706,6 +768,7 @@ def test_should_respond_403(self, unauthorized_test_client): "password": None, "schema": "new_schema", "description": TEST_CONN_DESCRIPTION, + "team_name": None, }, {"update_mask": ["schema", "extra"]}, ), @@ -818,6 +881,7 @@ def test_patch_should_respond_404(self, test_client, body): "password": "***", "port": 8080, "schema": None, + "team_name": None, }, {"update_mask": ["password"]}, ), @@ -833,6 +897,7 @@ def test_patch_should_respond_404(self, test_client, body): "password": "***", "port": 8080, "schema": None, + "team_name": None, }, {"update_mask": ["password"]}, ), @@ -853,6 +918,7 @@ def test_patch_should_respond_404(self, test_client, body): "password": "***", "port": 8080, "schema": None, + "team_name": None, }, {"update_mask": ["password", "extra"]}, ), diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py index 5e2c20780252f..28ae65fb01491 100644 --- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py +++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py @@ -203,6 +203,10 @@ class ConfigSection(BaseModel): options: Annotated[list[ConfigOption], Field(title="Options")] +class TeamName(RootModel[str]): + root: Annotated[str, Field(max_length=50, title="Team Name")] + + class ConnectionBody(BaseModel): """ Connection Serializer for requests body. @@ -220,6 +224,7 @@ class ConnectionBody(BaseModel): port: Annotated[int | None, Field(title="Port")] = None password: Annotated[str | None, Field(title="Password")] = None extra: Annotated[str | None, Field(title="Extra")] = None + team_name: Annotated[TeamName | None, Field(title="Team Name")] = None class ConnectionResponse(BaseModel): @@ -236,6 +241,7 @@ class ConnectionResponse(BaseModel): port: Annotated[int | None, Field(title="Port")] = None password: Annotated[str | None, Field(title="Password")] = None extra: Annotated[str | None, Field(title="Extra")] = None + team_name: Annotated[str | None, Field(title="Team Name")] = None class ConnectionTestResponse(BaseModel): @@ -914,10 +920,6 @@ class ValidationError(BaseModel): type: Annotated[str, Field(title="Error Type")] -class TeamName(RootModel[str]): - root: Annotated[str, Field(max_length=50, title="Team Name")] - - class VariableBody(BaseModel): """ Variable serializer for bodies.